Update app.py

app.py

@@ -23,15 +23,16 @@ 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(
-    torch.nn.Linear(in_ch, 1024),      # fc.1
-    torch.nn.ReLU(),
-    torch.nn.BatchNorm1d(1024),        # fc.3  
-    torch.nn.Linear(1024, 512),        # fc.5
-    torch.nn.ReLU(),
-    torch.nn.BatchNorm1d(512),         # fc.7
-    torch.nn.Linear(512, 196)          # fc.9
-)
+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)