Update app.py

app.py

@@ -63,7 +63,7 @@ def load_custom_model(model_key):
     # Pass model_name to config for correct model instantiation
     config = {"model": {"name": model_info["model_name"]}}
     model = XrayReg.load_from_checkpoint(model_info["ckpt"], map_location="cpu")
-    model = model.model
+    model = model.model.cuda() if torch.cuda.is_available() else model.model
     model.eval()
     for param in model.parameters():
         param.requires_grad = True
@@ -100,7 +100,7 @@ def predict_and_cam_custom(inp, model):
     # GradCAM for regression: use last conv layer, target output
     from pytorch_grad_cam import GradCAM
     from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
-    model = model.cuda()
+    # model = model.cuda()
     target_layers = [
         layer
         for name, layer in model.named_modules()
@@ -156,7 +156,7 @@ def download_models():
         local_dir=local_dir,
     )
 
-
+@spaces.GPU
 def main():
     # Download models if not already present
     try: