From dddb95caaca0d1ea873339207b0b37a1f371e263 Mon Sep 17 00:00:00 2001 From: Eric Silberstein Date: Wed, 19 Nov 2025 15:52:20 -0500 Subject: [PATCH] make mid_train script work even with a tiny number of iterations --- scripts/mid_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/mid_train.py b/scripts/mid_train.py index 6c2b82f..ebb6c42 100644 --- a/scripts/mid_train.py +++ b/scripts/mid_train.py @@ -139,7 +139,7 @@ def mid_data_generator(split): last_step = True # toggle last_step to True, which will terminate the training loop # Stopping condition to respect num_iterations, if given it += 1 - if num_iterations > 0 and it >= num_iterations: + if num_iterations > 0 and it >= num_iterations and split == "train": last_step = True # toggle last_step to True, which will terminate the training loop # Build up inputs/targets and yield for i in range(needed_tokens):