Merge pull request #1045 from bghira/main

merge
bghira · Oct 11, 2024 · 833c7f1 · 833c7f1
2 parents f608ef9 + 526d91e
commit 833c7f1
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/helpers/training/save_hooks.py b/helpers/training/save_hooks.py
@@ -491,11 +491,12 @@ def _load_full_model(self, models, input_dir):
 
     def load_model_hook(self, models, input_dir):
         # Check the checkpoint dir for a "training_state.json" file to load
-        if os.path.exists(self.training_state_path):
-            StateTracker.load_training_state(self.training_state_path)
+        training_state_path = os.path.join(input_dir, self.training_state_path)
+        if os.path.exists(training_state_path):
+            StateTracker.load_training_state(training_state_path)
         else:
             logger.warning(
-                f"Could not find {self.training_state_path} in checkpoint dir {input_dir}"
+                f"Could not find {training_state_path} in checkpoint dir {input_dir}"
             )
         if "lora" in self.args.model_type and self.args.lora_type == "standard":
             self._load_lora(models=models, input_dir=input_dir)