class_docstr: Subclass of a PyTorch `DataLoader` that will deal with device placement and current distributed setup.
  
  Args:
      dataset (`torch.utils.data.dataset.Dataset`):
          The dataset to use to build this datalaoder.
      device (`torch.device`, *optional*):
          If passed, the device to put all batches on.
      rng_types (list of `str` or [`~utils.RNGType`]):
          The list of random number generators to synchronize at the beginning of each iteration. Should be one or
          several of:
  
          - `"torch"`: the base torch random number generator
          - `"cuda"`: the CUDA random number generator (GPU only)
          - `"xla"`: the XLA random number generator (TPU only)
          - `"generator"`: an optional `torch.Generator`
      synchronized_generator (`torch.Generator`, *optional*):
          A random number generator to keep synchronized across processes.
      skip_batches (`int`, *optional*, defaults to 0):
          The number of batches to skip at the beginning.
      kwargs:
          All other keyword arguments to pass to the regular `DataLoader` initialization.
  
  **Available attributes:**
  
      - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
          Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
          number of processes
  
      - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.
class_docstr_tok: ["Subclass","of","a","PyTorch","`","DataLoader","`","that","will","deal","with","device","placement","and","current","distributed","setup",".","Args",":","dataset","(","`","torch.utils.data.dataset.Dataset","`",")",":","The","dataset","to","use","to","build","this","datalaoder",".","device","(","`","torch.device","`",",","*","optional","*",")",":","If","passed",",","the","device","to","put","all","batches","on",".","rng_types","(","list","of","`","str","`","or","[","`","~utils.RNGType","`","]",")",":","The","list","of","random","number","generators","to","synchronize","at","the","beginning","of","each","iteration",".","Should","be","one","or","several","of",":","-","`","``","torch","''","`",":","the","base","torch","random","number","generator","-","`","``","cuda","''","`",":","the","CUDA","random","number","generator","(","GPU","only",")","-","`","``","xla","''","`",":","the","XLA","random","number","generator","(","TPU","only",")","-","`","``","generator","''","`",":","an","optional","`","torch.Generator","`","synchronized_generator","(","`","torch.Generator","`",",","*","optional","*",")",":","A","random","number","generator","to","keep","synchronized","across","processes",".","skip_batches","(","`","int","`",",","*","optional","*",",","defaults","to","0",")",":","The","number","of","batches","to","skip","at","the","beginning",".","kwargs",":","All","other","keyword","arguments","to","pass","to","the","regular","`","DataLoader","`","initialization",".","*","*","Available","attributes",":","*","*","-","*","*","total_batch_size","*","*","(","`","int","`",")","--","Total","batch","size","of","the","dataloader","across","all","processes",".","Equal","to","the","original","batch","size","when","`","split_batches=True","`",";","otherwise","the","original","batch","size","*","the","total","number","of","processes","-","*","*","total_dataset_length","*","*","(","`","int","`",")","--","Total","length","of","the","inner","dataset","across","all","processes","."]
code: def total_batch_size(self):
      batch_sampler = (
          self.sampler
          if isinstance(self.sampler, BatchSampler)
          else self.batch_sampler
      )
      return (
          batch_sampler.batch_size
          if getattr(batch_sampler, "split_batches", False)
          else (
              batch_sampler.batch_size
              * getattr(batch_sampler, "num_processes", 1)
          )
      )
  
imports: import math
  from contextlib import suppress
  from typing import Callable, List, Optional, Union
  import torch
  from torch.utils.data import BatchSampler, DataLoader, IterableDataset, RandomSampler
  from .logging import get_logger
  from .state import AcceleratorState, DistributedType, GradientState, is_tpu_available
  from .utils import RNGType, broadcast, broadcast_object_list, concatenate, find_batch_size, get_data_structure, initialize_tensors, is_torch_version, send_to_device, slice_tensors, synchronize_rng_states
label_desc: Use image node_id 4 for calling the DataLoaderShard obj's underlying member method code with example usage: obj.total_batch_size() and returns: unknown