minor changes

train_tongue_diagnosis.py

@@ -87,7 +87,9 @@ def get_model(model_name):
         model.classifier = torch.nn.Linear(in_features=1024, out_features=NUM_CLASSES, bias=True)
         return model.to(device)
     elif model_name == 'resnet-50':
-        return models.resnet50(pretrained=True, num_classes=NUM_CLASSES).to(device)
+        model = models.resnet50(pretrained=True)
+        model.fc = torch.nn.Linear(in_features=2048, out_features=NUM_CLASSES, bias=True)
+        return model.to(device)
     elif model_name == 'vgg19_bn':
         model = models.vgg19_bn(pretrained=True)
         model.classifier[6] = torch.nn.Linear(in_features=4096, out_features=NUM_CLASSES, bias=True)