load to cpu first

app.py

@@ -394,17 +394,18 @@ def update_plot(dropdown, dropdown_choices_to_plot_args):
 
 
 # Set device
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-# device = torch.device('cuda')
+# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 # Load models
 models = [
     DDGPPIformer.load_from_checkpoint(
         PPIFORMER_WEIGHTS_DIR / f'ddg_regression/{i}.ckpt',
-        map_location=device
+        map_location=torch.device('cpu')
     ).eval()
     for i in range(3)
 ]
+models = [model.to(device) for model in models]
 
 
 app = gr.Blocks(theme=gr.themes.Default(primary_hue="green", secondary_hue="pink"))