initialize cuda

leo/inference.py

@@ -166,8 +166,8 @@ def form_batch(data_dict):
 
     
 def inference(scan_id, task, predict_mode=False):
-    # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    device = 'cpu' # ok for predict_mode=False, and both for Gradio demo local preview
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    # device = 'cpu' # ok for predict_mode=False, and both for Gradio demo local preview
 
     data_dict = load_data(scan_id)
     data_dict.update(get_lang(task))