update dataloader wrappers to have total_batch_size attribute

src/accelerate/data_loader.py

@@ -38,6 +38,15 @@
 if is_tpu_available(check_device=False):
     import torch_xla.distributed.parallel_loader as xpl
 
+    class MpDeviceLoaderWrapper(xpl.MpDeviceLoader):
+        def __init__(self, loader, device, **kwargs):
+            super().__init__(loader, device, **kwargs)
+
+        @property
+        def total_batch_size(self):
+            return self._loader.total_batch_size
+
+
 logger = get_logger(__name__)
 
 # kwargs of the DataLoader in min version 1.4.0.
@@ -321,6 +330,14 @@ def __iter__(self):
                 yield current_batch
                 break
 
+    @property
+    def total_batch_size(self):
+        return (
+            self.batch_sampler.batch_size
+            if self.batch_sampler.split_batches
+            else (self.batch_sampler.batch_size * self.batch_sampler.num_processes)
+        )
+
 
 class DataLoaderDispatcher(DataLoader):
     """
@@ -432,6 +449,12 @@ def __len__(self):
         else:
             return math.ceil(whole_length / self.state.num_processes)
 
+    @property
+    def total_batch_size(self):
+        return (
+            self.dataset.batch_size if self.split_batches else (self.dataset.batch_size * self.dataset.num_processes)
+        )
+
 
 def prepare_data_loader(
     dataloader: DataLoader,
@@ -577,7 +600,10 @@ def prepare_data_loader(
 
     if dispatch_batches:
         dataloader = DataLoaderDispatcher(
-            new_dataset, split_batches=split_batches, batch_sampler=new_batch_sampler, **kwargs
+            new_dataset,
+            split_batches=split_batches,
+            batch_sampler=new_batch_sampler,
+            **kwargs,
         )
     else:
         dataloader = DataLoaderShard(
@@ -590,5 +616,5 @@ def prepare_data_loader(
         )
 
     if state.distributed_type == DistributedType.TPU:
-        return xpl.MpDeviceLoader(dataloader, device)
+        return MpDeviceLoaderWrapper(dataloader, device)
     return dataloader