Does transformers 3.5.1 support auto mixed precision trainingļ¼Ÿ

Iā€™m using a GPTLMHead model in pytorch.
Is it possible , i add autocast() in the forward function in GPTLMHead and change the training process followed the Automatic Mixed Precision ā€” PyTorch Tutorials 1.8.1+cu102 documentation

Yes, itā€™s possible.

I train the gpt model using V100 with mixed precision training. But it happens that in the inference time, the generated token is mess with the autocast() open.when i turn off it, everything is ok.