Mixed precision on part of the model

Hi, everyone. Can I do mixed precision training only on the part of the model? Suppose, I have a model, which has a backbone and a classifier. I want to train backbone with mixed precision and classifier with single precision. Is there any way to do it?

1 Like

Yes, you can wrap only certain parts of the model into with torch.cuda.amp.autocast() or use the decorator on specific functions.