Hi.
I am currently trying to figure out how to facilitate mixed-precision training when using transfer learning.
Typically, here’s how I am building my custom model using a pre-trained base model:
# load up the ResNet50 model
model = models.resnet50(pretrained=True)

# append a new classification top to our feature extractor and pop it
# on to the current device
modelOutputFeats = model.fc.in_features
model.fc = nn.Linear(modelOutputFeats, len(trainDS.classes))
model = model.to(DEVICE)

I am also using nn.DataParallel. As demonstrated here 1, simply putting model inside autocast context, won’t help. So, is there a way I could perform it using the practices recommended here?