Move to device inside @spaces.GPU

app.py

@@ -70,6 +70,9 @@ random.seed(0)
 
 @spaces.GPU
 def predict_ddg(models, ppi, muts, return_attn):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f"[INFO] Device on prediction: {device}")
+    models = [model.to(device) for model in models]
     if return_attn:
         ddg_pred, attns = predict_ddg_(models, ppi, muts, return_attn=return_attn)
         return ddg_pred.detach().cpu(), attns.detach().cpu()
@@ -503,8 +506,8 @@ with app:
     download_from_zenodo('weights.zip')
 
     # Set device
-    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f"[INFO] Device on start: {device}")
 
     # Load models
     models = [