Sayak_Paul:

My question is how do I achieve that when I am using a pre-trained model?


You could write a custom model, decorate the forward method of the custom model with the autocast context, and call the pretrained model inside.