Tutel with pytorch automatic mixed precision package

[MiZhenxing]

Hi, thank you for your excellent package. I wonder whether tutel could be used seamless together with the automatic mixed precision package of PyTorch. If so, could you provide some hints on how to use it or provide an example python script?
Thank you very much!

[ghostplant]

Yes, Tutel is able to support native Pytorch AMP. Please follow this example: https://github.com/microsoft/tutel/blob/main/tutel/examples/helloworld_amp.py#L76

Note that you need to use @autocast() and with autocast: properly according to pytorch's doc.
Welcome to report problems if you encounter any issue.

[MiZhenxing]

Thanks for your quick response. Sorry I didn't notice that example 🤣. I will do some testing on my code.