Actually reverse states

huggingface · Oct 10, 2024 · 49e6556 · 49e6556
1 parent 8499302
commit 49e6556
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
@@ -268,7 +268,7 @@ def load_accelerator_state(
 
     # Random states
     try:
-        states = load(input_dir.joinpath(f"{RNG_STATE_NAME}_{process_index}.pkl"), weights_only=False)
+        states = load(input_dir.joinpath(f"{RNG_STATE_NAME}_{process_index}.pkl"), weights_only=True)
         if "step" in states:
             override_attributes["step"] = states["step"]
         random.setstate(states["random_state"])