Update model_utils.py

model_utils.py

@@ -24,13 +24,16 @@ class CheXNet(nn.Module):
         x = torch.flatten(x, 1)
         return self.classifier(x)
 
-def load_model(url, device):
-    model_path = "dannynet.pth"
-    torch.hub.download_url_to_file(url,model_path)
-    model = torch.load(model_path, map_location = device)
+from huggingface_hub import hf_hub_download
+import torch
+
+def load_model(repo_id, filename, device):
+    model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    model = torch.load(model_path, map_location=device)
     model.eval()
     return model
 
+
 def predict(model, img_tensor, device):
     with torch.no_grad():
         output = model(img_tensor.unsqueeze(0).to(device))