From f51ab0383d86c73f4e3ceacd8716a615b2ffe9ee Mon Sep 17 00:00:00 2001
From: Igor Shilov <hello@igorshilov.com>
Date: Thu, 3 Jul 2025 17:56:28 +0000
Subject: [PATCH] fix np.int32 overflow in token counting

---
 train_gpt2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train_gpt2.py b/train_gpt2.py
index b9dee87..5538475 100644
--- a/train_gpt2.py
+++ b/train_gpt2.py
@@ -342,7 +342,7 @@ class DistributedDataLoader:
         for fname in self.files:
             shard_ntok = _peek_data_shard(fname)
             assert shard_ntok >= num_processes * B * T + 1
-            ntok_total += shard_ntok
+            ntok_total += int(shard_ntok)
         self.ntok_total = ntok_total
         print0(f"DataLoader: total number of tokens: {ntok_total:,} across {len(self.files)} files")