Hotfix for final step model saving

simple_stories_train/train_llama.py

@@ -456,7 +456,7 @@ def get_lr(it: int) -> float:
             and master_process
             and (
                 (config.intermediate_checkpoints and is_checkpoint_step(step))
-                or last_step
+                or step == config.num_iterations - 1
                 or train_loader_depleted
             )
         ):