Update nanochat/checkpoint_manager.py

Co-authored-by: Sofie Van Landeghem <svlandeg@users.noreply.github.com>
This commit is contained in:
Pyry Takala 2025-11-21 13:47:22 -08:00 committed by GitHub
parent 3e2a0668b2
commit 958543fcd5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -117,7 +117,7 @@ def find_last_step(checkpoint_dir):
checkpoint_files = [f for f in os.listdir(checkpoint_dir) if re.search(r'model_(\d+)\.pt$', f)] checkpoint_files = [f for f in os.listdir(checkpoint_dir) if re.search(r'model_(\d+)\.pt$', f)]
if not checkpoint_files: if not checkpoint_files:
raise FileNotFoundError(f"No checkpoints found in {checkpoint_dir}") raise FileNotFoundError(f"No checkpoints found in {checkpoint_dir}")
last_step = int(max(re.search(r'model_(\d+)\.pt$', f).group(1) for f in checkpoint_files)) last_step = int(max(f.split("_")[-1].split(".")[0] for f in checkpoint_files))
return last_step return last_step
# ----------------------------------------------------------------------------- # -----------------------------------------------------------------------------