fix bugs in model restore.

chenyaofo · Apr 11, 2021 · 948fc63 · 948fc63
1 parent 19f1077
commit 948fc63
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/codebase/torchutils/common.py b/codebase/torchutils/common.py
@@ -331,7 +331,8 @@ def save(self, minitor: str, metrics: dict, states: dict):
     def restore(self, metrics: dict, states: dict, device="cuda:0"):
         checkpoint_path = self.output_directory / "checkpoint.pt"
         if checkpoint_path.exists():
-            checkpoint: dict = torch.load(checkpoint_path, map_location=device)
+            map_location= f"cuda:{device}" if isinstance(device, int) else device
+            checkpoint: dict = torch.load(checkpoint_path, map_location=map_location)
             metrics.update(checkpoint.pop("metrics", dict()))
             for name, module in states.items():
                 module.load_state_dict(checkpoint[name])