Update model_helper.py

model_helper.py

@@ -104,7 +104,7 @@ def load_model_checkpoint(args=None, device='cpu'):
     print(f"Task: {tm.task_name}, Max Shift Steps: {tm.max_shift_steps}")
 
     # Use GPU if available
-    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     # Model
     model = YourMT3(