detach predictions

app.py

@@ -60,8 +60,13 @@ random.seed(0)
 
 
 @spaces.GPU
-def predict_ddg(*args, **kwargs):
-    return predict_ddg_(*args, **kwargs)
+def predict_ddg(models, ppi, muts, return_attn):
+    if return_attn:
+        ddg_pred, attns = predict_ddg_(models, ppi, muts, return_attn=return_attn)
+        return ddg_pred.detach().cpu(), attns.detach().cpu()
+    else:
+        ddg_pred = predict_ddg_(models, ppi, muts, return_attn=return_attn)
+        return ddg_pred.detach().cpu()
 
 
 def process_inputs(inputs, temp_dir):