Update model_utils.py

model_utils.py

@@ -24,16 +24,26 @@ class CheXNet(nn.Module):
         x = torch.flatten(x, 1)
         return self.classifier(x)
 
-from huggingface_hub import hf_hub_download
-import torch
 
+
+from huggingface_hub import hf_hub_download
 def load_model(device):
     model_path = hf_hub_download(repo_id="cfgpp/danny_net", filename="dannynet.pth")
-    model = torch.load(model_path, map_location=device)
+    
+    # Rebuild model architecture
+    model = CheXNet(num_classes=14)
+    
+    # Load state dict (just weights)
+    state_dict = torch.load(model_path, map_location=device)
+    model.load_state_dict(state_dict)
+    
+    model.to(device)
     model.eval()
     return model
 
 
+
+
 def predict(model, img_tensor, device):
     with torch.no_grad():
         output = model(img_tensor.unsqueeze(0).to(device))