add total samples

PrimeIntellect-ai · Sep 21, 2024 · 5d0deb9 · 5d0deb9
1 parent 9a5efa7
commit 5d0deb9
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/zeroband/train.py b/src/zeroband/train.py
@@ -195,6 +195,7 @@ def train(config: Config):
                 * config.optim.batch_size
                 / (time.time() - beginning_step_time),
                 "Perplexity": torch.exp(loss_batch).item(),
+                "total_tokens": real_step * config.optim.batch_size * config.data.seq_length,
             }
 
             if world_info.rank == 0: