diff --git a/train_arabic_r8_hf.py b/train_arabic_r8_hf.py
index 299e095..407f4b0 100644
--- a/train_arabic_r8_hf.py
+++ b/train_arabic_r8_hf.py
@@ -173,6 +173,9 @@ def add_silver(tgt: str) -> None:
     tok = AutoTokenizer.from_pretrained(TEACHER)
 
     class DS(Dataset):
+        def __len__(self):
+            return len(units)
+
         def __getitem__(self, i):
             src, tgt = units[i]
             e = tok(src, truncation=True, max_length=WINDOW)