checkpoint hot fix (#140)

allenai · Jan 17, 2025 · 48abe8c · 48abe8c
1 parent 0c096e2
commit 48abe8c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/olmo_core/train/checkpoint.py b/src/olmo_core/train/checkpoint.py
@@ -171,6 +171,7 @@ def load(
             for path in (f"{dir}/train/rank{get_rank()}.pt", f"{dir}/train/rank0.pt"):
                 try:
                     trainer_state = torch.load(cached_path(path, quiet=True), weights_only=False)
+                    break
                 except FileNotFoundError:
                     pass