Update app.py

app.py

@@ -24,16 +24,17 @@ model = resnet50(weights=None)
 
 # Replace fc with the multi-layer head that matches checkpoint
 in_ch = model.fc.in_features
-model.fc = torch.nn.Sequential()
-model.fc.add_module('1', torch.nn.Linear(in_ch, 1024))      # fc.1
-model.fc.add_module('2', torch.nn.ReLU())
-model.fc.add_module('3', torch.nn.BatchNorm1d(1024))        # fc.3  
-model.fc.add_module('4', torch.nn.ReLU())
-model.fc.add_module('5', torch.nn.Linear(1024, 512))        # fc.5
-model.fc.add_module('6', torch.nn.ReLU())
-model.fc.add_module('7', torch.nn.BatchNorm1d(512))         # fc.7
-model.fc.add_module('8', torch.nn.ReLU())
-model.fc.add_module('9', torch.nn.Linear(512, 196))         # fc.9
+model.fc = torch.nn.Linear(in_ch, num_classes)
+# model.fc = torch.nn.Sequential()
+# model.fc.add_module('1', torch.nn.Linear(in_ch, 1024))      # fc.1
+# model.fc.add_module('2', torch.nn.ReLU())
+# model.fc.add_module('3', torch.nn.BatchNorm1d(1024))        # fc.3  
+# model.fc.add_module('4', torch.nn.ReLU())
+# model.fc.add_module('5', torch.nn.Linear(1024, 512))        # fc.5
+# model.fc.add_module('6', torch.nn.ReLU())
+# model.fc.add_module('7', torch.nn.BatchNorm1d(512))         # fc.7
+# model.fc.add_module('8', torch.nn.ReLU())
+# model.fc.add_module('9', torch.nn.Linear(512, 196))         # fc.9
 
 # Load the state dict
 model.load_state_dict(state_dict, strict=True)