Update app.py

app.py

@@ -23,19 +23,24 @@ def download_model():
 # Load the model from Hugging Face
 def load_model(model_path):
     model = models.resnet50(pretrained=False)
-    model.fc = nn.Linear(model.fc.in_features, 3)  # Adjust for 3 classes
+    num_features = model.fc.in_features
+    model.fc = nn.Sequential(
+        nn.Dropout(0.5),
+        nn.Linear(num_features, 3)  # 3 classes
+    )
     checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
-    
-    # Check if it is a checkpoint and extract model state dict
-    if 'model_state_dict' in checkpoint:
-        state_dict = checkpoint['model_state_dict']
-    else:
-        state_dict = checkpoint
-    
-    model.load_state_dict(state_dict)
+    model.load_state_dict(checkpoint['model_state_dict'])
+
+    # Rename keys to match the model definition
+state_dict['fc.weight'] = state_dict.pop('fc.1.weight')
+state_dict['fc.bias'] = state_dict.pop('fc.1.bias')
+
+# Load the modified state dict
+model.load_state_dict(state_dict)
     model.eval()
     return model
 
+
 # Path to your model
 model_path = hf_hub_download(repo_id="jays009/Resnet3", filename="model.pth")
 model = load_model(model_path)