Update app.py

app.py

@@ -84,13 +84,12 @@ def predict_custom(model, input_tensor):
     return float(pred)
 
 
-@spaces.GPU
 def load_custom_model(model_key):
     model_info = MODEL_OPTIONS[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.cuda() if torch.cuda.is_available() else model.model
+    model = model.model
     model.eval()
     for param in model.parameters():
         param.requires_grad = True
@@ -99,6 +98,10 @@ def load_custom_model(model_key):
 @spaces.GPU
 def predict_and_cam_custom(inp, model):
     input_tensor, rgb_img = preprocess_image_custom(inp)
+    model = model.cuda() if torch.cuda.is_available() else model.model
+    model.eval()
+    for param in model.parameters():
+        param.requires_grad = True
     value = predict_custom(model, input_tensor)
     # GradCAM for regression: use last conv layer, target output
     from pytorch_grad_cam import GradCAM