Bugfix.

mlfoundations · May 13, 2024 · feee5b6 · feee5b6
1 parent f01a34e
commit feee5b6
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/open_lm/main.py b/open_lm/main.py
@@ -824,7 +824,12 @@ def main(args):
             logging.info("Training exiting due to NaN value")
             break
 
-        expected_steps = sum(num_samples_per_source) // args.global_batch_size
+        if args.dataset_manifest is not None:
+            expected_samples = sum(num_samples_per_source)
+        else:
+            expected_samples = args.train_num_samples
+
+        expected_steps = expected_samples // args.global_batch_size
         if steps_done_epoch < (1 - args.data_tolerate_error_p) * expected_steps:
             num_ckpt_too_few_tokens += 1