last changes

Helsinki-NLP · Dec 11, 2024 · 038c0a6 · 038c0a6
1 parent 75d2670
commit 038c0a6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mammoth/utils/model_saver.py b/mammoth/utils/model_saver.py
@@ -88,7 +88,7 @@ def load_parameters_from_checkpoint(
         name = component.get_name()
         checkpoint_path = f'{checkpoint_prefix}_{name}.pt'
         if os.path.isfile(checkpoint_path):
-            state_dict = torch.load(checkpoint_path)
+            state_dict = torch.load(checkpoint_path, map_location="cpu")
             incompatible_keys = component.load_state_dict(model=model, state_dict=state_dict)
             if incompatible_keys.missing_keys or incompatible_keys.unexpected_keys:
                 logger.info(f'Module {name} incompatible keys: {incompatible_keys}')
@@ -103,7 +103,7 @@ def load_parameters_from_checkpoint(
             optimizer_path = f'{checkpoint_prefix}_{name}_optim.pt'
             if os.path.isfile(optimizer_path):
                 # The optimizer parameters are distributed the same way as the components
-                optim_state_dict = torch.load(optimizer_path)
+                optim_state_dict = torch.load(optimizer_path, map_location="cpu")
                 incompatible_keys = optim.suboptimizers[name].load_state_dict(optim_state_dict)
                 if incompatible_keys and (incompatible_keys.missing_keys or incompatible_keys.unexpected_keys):
                     logger.info(f'Optim {name} incompatible keys: {incompatible_keys}')