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?
Yes, you can wrap only certain parts of the model into
with torch.cuda.amp.autocast() or use the decorator on specific functions.