Explicit check for step when loading the state (#2992)

* Explicit check * Nit
huggingface · Aug 6, 2024 · d982751 · d982751
1 parent 95edc68
commit d982751
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -3159,7 +3159,8 @@ def _inner(folder):
             map_location,
             **load_model_func_kwargs,
         )
-        self.step = override_attributes["step"]
+        if "step" in override_attributes:
+            self.step = override_attributes["step"]
         custom_checkpoints = [
             f for f in os.listdir(input_dir) if re.search(r"^custom_checkpoint_\d+\.pkl$", f) is not None
         ]

diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
@@ -252,7 +252,8 @@ def load_accelerator_state(
     # Random states
     try:
         states = torch.load(input_dir.joinpath(f"{RNG_STATE_NAME}_{process_index}.pkl"))
-        override_attributes["step"] = states["step"]
+        if "step" in states:
+            override_attributes["step"] = states["step"]
         random.setstate(states["random_state"])
         np.random.set_state(states["numpy_random_seed"])
         torch.set_rng_state(states["torch_manual_seed"])