Solved by vgsprasad in post #3 
              
              
                Define your model as follows. 
from torchvision import models 
class FinetuneResnet(nn.Module):
    def __init__(self, num_classes):
        super(FinetuneResnet, self).__init__()

        self.model = models.resnet50(pretrained=True)
        self.fc1 = nn.Linear(2048, 2048)
        self.fc2 = nn.Li…