[BUG]: Shardformer failure with torch 2.3

[Edenzzzz]

Is there an existing issue for this bug?

• I have searched the existing issues

🐛 Describe the bug

Installing Pytorch 2.3, which is required by the newest version of xformers, can cause shardformer to fail for obscure reasons. This issue is opened for tracking purposes and to discourage the use of torch 2.3 for now.

To reproduce:

python tests/test_optimizer/test_dist_lamb.py

Environment

PyTorch 2.3