Fix length of DispatchDataLoader (#175)

huggingface · Sep 27, 2021 · 929e17d · 929e17d
1 parent d6247b7
commit 929e17d
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 from typing import List, Optional, Union
 
 import torch
@@ -388,6 +389,14 @@ def __iter__(self):
                 xm.mark_step()
             yield slice_tensors(batch, data_slice)
 
+    def __len__(self):
+        state = AcceleratorState()
+        whole_length = super().__len__()
+        if self.drop_last:
+            return whole_length // state.num_processes
+        else:
+            return math.ceil(whole_length / state.num_processes)
+
 
 def prepare_data_loader(
     dataloader: DataLoader,