Skip to content

Commit

Permalink
minor logging improvement
Browse files Browse the repository at this point in the history
  • Loading branch information
epwalsh committed Jan 22, 2025
1 parent 212108f commit 7b755c9
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions src/olmo_core/train/checkpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -193,14 +193,15 @@ def load(
if get_rank(self.process_group) == 0:
try:
metadata = get_checkpoint_metadata(model_and_optim_dir)
except FileNotFoundError:
except FileNotFoundError as exc:
# Try base directory, which could be the case if user is trying to load model weights
# (possibly with optimizer state), and not an actual train checkpoint.
if trainer_state is None:
metadata = get_checkpoint_metadata(dir)
model_and_optim_dir = dir
else:
raise
raise FileNotFoundError(f"Missing checkpointing metadata in '{dir}'") from exc

if load_optimizer_state is None:
for key in metadata.state_dict_metadata.keys():
if key.startswith("optim."):
Expand Down

0 comments on commit 7b755c9

Please sign in to comment.