Let iterable dataset shard have a len (#2066)

huggingface · Oct 23, 2023 · 11e2e99 · 11e2e99
1 parent 07e745f
commit 11e2e99
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
@@ -271,6 +271,13 @@ def __init__(
         self.process_index = process_index
         self.split_batches = split_batches
 
+    def __len__(self):
+        # We will just raise the downstream error if the underlying dataset is not sized
+        if self.drop_last:
+            return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size
+        else:
+            return math.ceil(len(self.dataset) / (self.batch_size * self.num_processes)) * self.batch_size
+
     def __iter__(self):
         real_batch_size = self.batch_size if self.split_batches else (self.batch_size * self.num_processes)
         process_batch_size = (self.batch_size // self.num_processes) if self.split_batches else self.batch_size