Update app.py

app.py

@@ -176,8 +176,12 @@ multiclass_model_path = hf_hub_download(repo_id="santhoshraghu/DermBOT", filenam
 multilabel_model = SkinViT(num_classes=len(multilabel_class_names))
 multiclass_model = DermNetViT(num_classes=len(multiclass_class_names))
 
-multilabel_model.load_state_dict(torch.load(multilabel_model_path, map_location="cpu"))
-multiclass_model.load_state_dict(torch.load(multiclass_model_path, map_location="cpu"))
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+multilabel_model.load_state_dict(torch.load(multilabel_model_path, map_location=device))
+multiclass_model.load_state_dict(torch.load(multiclass_model_path, map_location=device))
+multilabel_model.to(device)
+multiclass_model.to(device)
+
 
 multilabel_model.eval()
 multiclass_model.eval()
@@ -193,7 +197,7 @@ def run_inference(image):
         transforms.ToTensor(),
         transforms.Normalize([0.5], [0.5])
     ])
-    input_tensor = transform(image).unsqueeze(0)
+    input_tensor = transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
         probs_multi = torch.sigmoid(multilabel_model(input_tensor)).squeeze().numpy()
         predicted_multi = [multilabel_class_names[i] for i, p in enumerate(probs_multi) if p > 0.5]