Run on GPU if available, otherwise CPU

app.py

@@ -33,7 +33,8 @@ def load_model():
         token=os.getenv("HF_TOKEN")
     )
 
-    model = torch.load(ckpt_path, map_location="cuda")  # no weights_only, no globals hack
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = torch.load(ckpt_path, map_location=device)
     model = disable_dropout(model)
     model.to("cuda")
     model.eval()