Fix model saving at end of training

danbraunai · Dec 1, 2024 · 476acdd · 476acdd
1 parent 1fee6ab
commit 476acdd
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/simple_stories_train/train_llama.py b/simple_stories_train/train_llama.py
@@ -454,7 +454,11 @@ def get_lr(it: int) -> float:
         if (
             checkpoints_dir is not None
             and master_process
-            and (config.intermediate_checkpoints and is_checkpoint_step(step) or last_step)
+            and (
+                (config.intermediate_checkpoints and is_checkpoint_step(step))
+                or last_step
+                or train_loader_depleted
+            )
         ):
             save_model(checkpoints_dir, raw_model, step=step, wandb_project=config.wandb_project)