Update app.py

app.py

@@ -23,12 +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, num_classes)
-    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+    model.fc = nn.Linear(model.fc.in_features, 3)  # Adjust for 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.eval()
-    logging.info("Model loaded successfully. Ready for inference.")
     return model
 
+# Path to your model
+model_path = hf_hub_download(repo_id="jays009/Resnet3", filename="model.pth")
+model = load_model(model_path)
+
+
 # Download the model and load it
 model_path = download_model()
 model = load_model(model_path)