fixed load peft weight. #47

shibing624 · Jun 20, 2023 · adf0335 · adf0335
1 parent 9e7c163
commit adf0335
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/textgen/chatglm/chatglm_model.py b/textgen/chatglm/chatglm_model.py
@@ -327,7 +327,7 @@ def train_model(
                 # The two files above have a different name depending on how they were saved, but are actually the same.
                 if os.path.exists(checkpoint_name):
                     logger.info(f"Restarting from {checkpoint_name}")
-                    adapters_weights = torch.load(checkpoint_name)
+                    adapters_weights = torch.load(checkpoint_name, map_location='cpu')
                     set_peft_model_state_dict(self.model, adapters_weights)
                 else:
                     logger.warning(f"Checkpoint {checkpoint_name} not found")

diff --git a/textgen/gpt/gpt_model.py b/textgen/gpt/gpt_model.py
@@ -356,7 +356,7 @@ def train_model(
                 # The two files above have a different name depending on how they were saved, but are actually the same.
                 if os.path.exists(checkpoint_name):
                     logger.info(f"Restarting from {checkpoint_name}")
-                    adapters_weights = torch.load(checkpoint_name)
+                    adapters_weights = torch.load(checkpoint_name, map_location='cpu')
                     set_peft_model_state_dict(self.model, adapters_weights)
                 else:
                     logger.warning(f"Checkpoint {checkpoint_name} not found")