So, here’s what I did:
from torch.cuda.amp import autocast
from torch import nn

class CustomClassifier(nn.Module):
	def __init__(self, baseModel, numClasses):
		super(FoodClassier, self).__init__()

		# initialize the base model and the classification layer
		self.baseModel = baseModel
		self.classifier = nn.Linear(baseModel.classifier.in_features,
			numClasses)

		# set the classifier of our base model to produce outputs
		# from the last convolution block
		self.baseModel.classifier = nn.Identity()

	# we decorate the `forward()` method with `autocast()` to enable
	# mixed-precision training in a distributed manner
	@autocast(enabled=True)
	def forward(self, x):
		# pass the inputs through the base model and then obtain the
		# obtain the classifier outputs
		features = self.baseModel(x)
		logits = self.classifier(features)

		# return the classifier outputs
		return logits

I am using a densenet121 hence changed it to classifier from fc. However, I don’t see any speed-up in the total model training after making these changes.