Update app.py

app.py

@@ -58,16 +58,7 @@ def preprocess_image(inp):
         raise
 
 
-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.eval()
-    for param in model.parameters():
-        param.requires_grad = True
-    return model
+
 
 
 def preprocess_image_custom(inp):
@@ -93,6 +84,18 @@ 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.eval()
+    for param in model.parameters():
+        param.requires_grad = True
+    return model
+
 @spaces.GPU
 def predict_and_cam_custom(inp, model):
     input_tensor, rgb_img = preprocess_image_custom(inp)