Fix MLM datasets length

thomasw21 · thomasw21 · commit 9da75889ac0c · 2022-06-27T20:48:54.000+02:00
diff --git a/megatron/data/mlm_dataset.py b/megatron/data/mlm_dataset.py
@@ -335,7 +335,7 @@ def __init__(
             assert self.inputs_length + self.targets_length == args.seq_length
 
     def __len__(self):
-        return len(self.samples_mapping)
+        return len(self._gpt_dataset)
 
     def __getitem__(self, idx):
         if isinstance(idx, slice):
