Fix length of IterableDatasetShard and add test (huggingface#13792)

* Fix length of IterableDatasetShard and add test * Add comments
Albertobegue · Jan 27, 2022 · 3af7ddd · 3af7ddd
1 parent a220407
commit 3af7ddd
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/src/transformers/trainer_pt_utils.py b/src/transformers/trainer_pt_utils.py
@@ -775,9 +775,9 @@ def __iter__(self):
     def __len__(self):
         # Will raise an error if the underlying dataset is not sized.
         if self.drop_last:
-            return len(self.dataset) // self.num_processes
+            return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size
         else:
-            return math.ceil(len(self.dataset) / self.num_processes)
+            return math.ceil(len(self.dataset) / (self.batch_size * self.num_processes)) * self.batch_size
 
 
 # In order to keep `trainer.py` compact and easy to understand, place any secondary PT Trainer

diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py
@@ -281,6 +281,7 @@ def create_reverse_dependency_map():
         "test_trainer_distributed.py",
         "test_trainer_tpu.py",
     ],
+    "train_pt_utils.py": "test_trainer_utils.py",
     "utils/versions.py": "test_versions_utils.py",
 }