Bugfix: precise device in torch.load

ThomasRieutord · Dec 5, 2023 · 1e49f4f · 1e49f4f
1 parent 56cf264
commit 1e49f4f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/mmt/inference/io.py b/mmt/inference/io.py
@@ -199,7 +199,7 @@ def load_old_pytorch_model(xp_name, lc_in="esawc", lc_out="esgp"):
 
     return model
 
-def load_pytorch_model(xp_name, lc_in="esawc", lc_out="esgp", train_mode = False):
+def load_pytorch_model(xp_name, lc_in="esawc", lc_out="esgp", train_mode = False, device = "cpu"):
     """Return the pre-trained Pytorch model from the experiment `xp_name`"""
 
     if os.path.isabs(xp_name):
@@ -223,7 +223,7 @@ def load_pytorch_model(xp_name, lc_in="esawc", lc_out="esgp", train_mode = False
 
     assert os.path.isfile(checkpoint_path), f"No checkpoint found at {checkpoint_path}"#
 
-    checkpoint = torch.load(checkpoint_path)
+    checkpoint = torch.load(checkpoint_path, map_location=device)
     config = utilconf.get_config(config_path)
 
     if config.model.type == "transformer_embedding":

diff --git a/mmt/inference/translators.py b/mmt/inference/translators.py
@@ -279,7 +279,7 @@ def __init__(
             self.esawc.n_labels + 1, device=self.device
         )
         self.encoder_decoder = io.load_pytorch_model(
-            checkpoint_path, lc_in="esawc", lc_out="esgp"
+            checkpoint_path, lc_in="esawc", lc_out="esgp", device = device
         )
         self.encoder_decoder.to(self.device)
         self.landcover = self.esawc