Update main.py

zjunlp · Jul 28, 2024 · a722538 · a722538
1 parent ab28122
commit a722538
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/MarT/main.py b/MarT/main.py
@@ -156,7 +156,7 @@ def load_state_dict():
         trainer.fit(lit_model, datamodule=data)
         path = model_checkpoint.best_model_path
         # load best model
-        lit_model.load_state_dict(torch.load(path, map_location='cuda:2')["state_dict"])
+        lit_model.load_state_dict(torch.load(path, map_location='cuda')["state_dict"])
 
     result = trainer.test(lit_model, datamodule=data)
     print(result)