Hi Alex, I went through your code in the model.py file. Here’s what I thought can be changed to make it train in a better way.
Typically, when doing Transfer Learning, we use take a network trained on a dataset D, remove the final classification layer used for that dataset D, and then add our own final layer for classification. The objective here is to learn from the final encodings learnt from dataset D. Hence, a better way of utilizing the pre-trained network is shown below for VGG11.
""" Custom VGG11 """ 
class VGG11(nn.Module): 
    def __init__(self): 
        super(VGG11, self).__init__()
        self.net = models.vgg11(pretrained=True)
        # Replacing final classification layer with custom classification layer
        self.net.classifier[-1] = nn.Linear(4096, 2)
        self.net.cuda().half() 
    
    def forward(self, x):
        f = self.net(x) 
        return f