Fix error in final ckpt

mlfoundations · May 13, 2024 · afee0c7 · afee0c7
1 parent 543ae63
commit afee0c7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/open_lm/main.py b/open_lm/main.py
@@ -825,7 +825,7 @@ def main(args):
             break
 
         expected_steps = data["train"].dataloader.total_batches
-        if steps_done_epoch < (1 - args.data_tolerate_error_p) * expected_steps:
+        if steps_done_epoch < (1 - args.data_tolerate_error_p) * expected_steps and not done_training:
             num_ckpt_too_few_tokens += 1
 
         if num_ckpt_too_few_tokens > args.data_tolerate_num_ckpts: