@ptrblck the linked doc (first one) discusses AMP when using nn.DataParallel. It states that the forward method of my model needs to be inside an autocast context.
My question is how do I achieve that when I am using a pre-trained model? Would the snippet I provided work in that case?
The autocast state would refer to the with torch.cuda.amp.autocast(enabled=True): context.