From 5d0deb9c2c3852e80957004628af96bfc71628e3 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Sat, 21 Sep 2024 00:19:50 +0000
Subject: [PATCH] add total samples

---
 src/zeroband/train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/zeroband/train.py b/src/zeroband/train.py
index 9efc8367..056d89ac 100644
--- 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: