def __init__(self, dataset_dict, length=None):
        """dataset_dict: dictionary mapping from index to batch
        length is used in the case of DistributedSampler: e.g. the dataset could have size 1k, but
        with 8 GPUs the dataset_dict would only have 125 items.
        """
        super().__init__()
        self.dataset_dict = dataset_dict
        self.length = length or len(self.dataset_dict)
dataset_dict: dictionary mapping from index to batch
        length is used in the case of DistributedSampler: e.g. the dataset could have size 1k, but
        with 8 GPUs the dataset_dict would only have 125 items.
        
training/src/datamodules/imagenet.py