## What does this PR do?
Fixes # (issue).

## Before submitting

- [ ] Did you have fun?
  - Make sure you had fun coding 🙃
- [ ] Did you read the [contributor guideline](https://github.com/facebookresearch/fairscale/blob/main/CONTRIBUTING.md)?
- [ ] Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  - [ ] N/A
- [ ] Did you make sure to update the docs?
  - [ ] N/A
- [ ] Did you write any new necessary tests?
  - [ ] N/A
- [ ] Did you update the [changelog](https://github.com/facebookresearch/fairscale/blob/main/CHANGELOG.md)? (if needed)
  - [ ] N/A


## PR review
Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.
