This doesn't seem to build with torch 1.8

[min-xu-ai]

See this:

https://app.circleci.com/pipelines/github/facebookresearch/fairscale/1431/workflows/3b826aea-7ccc-4ca9-b817-2788488a85af/jobs/6482

[min-xu-ai]

I will try to fix it and send a PR