Solved by yash1994 in post #4 
              
              
                Check the implementation of ResNet50 from torchvision and implement custom network something like: 
class ModifiedResNet50(nn.Module):
    def __init__(self,):
         pass
    def forward(input):
         pass

Then load weights from pretrained ResNet50 like given below: 
vanilla_resnet = models.r…