def __init__(
        self,
        data_dir: str,
        image_size: int = 224,
        train_transforms=None,
        val_transforms=None,
        test_transforms=None,
        img_dtype='float32',  # Using str since OmegaConf doesn't support non-primitive type
        cache_val_dataset=False,
        mixup: Optional[Callable] = None,
        num_aug_repeats: int = 0,
        num_workers: int = 0,
        batch_size: int = 32,
        batch_size_eval: Optional[int] = None,
        shuffle: bool = True,
        pin_memory: bool = True,
        drop_last: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        """
        Args:
            data_dir: path to the imagenet dataset file
            num_imgs_per_val_class: how many images per class for the validation set
            image_size: final image size
            num_workers: how many data workers
            batch_size: batch_size
            shuffle: If true shuffles the data every epoch
            pin_memory: If true, the data loader will copy Tensors into CUDA pinned memory before
                        returning them
            drop_last: If true drops the last incomplete batch
        """
        super().__init__(*args, **kwargs)

        self.image_size = image_size
        self.train_transforms = train_transforms
        self.val_transforms = val_transforms
        self.test_transforms = test_transforms
        assert img_dtype in ['float32', 'float16', 'bfloat16']
        self.img_dtype = torch.__getattribute__(img_dtype)
        self.cache_val_dataset = cache_val_dataset
        self.mixup = mixup
        self.num_aug_repeats = num_aug_repeats
        self.dims = (3, self.image_size, self.image_size)
        self.data_dir = Path(data_dir).expanduser()
        self.num_workers = num_workers
        self.batch_size = batch_size
        self.batch_size_eval = batch_size_eval if batch_size_eval is not None else self.batch_size
        self.shuffle = shuffle
        self.pin_memory = pin_memory
        self.drop_last = drop_last

        Args:
            data_dir: path to the imagenet dataset file
            num_imgs_per_val_class: how many images per class for the validation set
            image_size: final image size
            num_workers: how many data workers
            batch_size: batch_size
            shuffle: If true shuffles the data every epoch
            pin_memory: If true, the data loader will copy Tensors into CUDA pinned memory before
                        returning them
            drop_last: If true drops the last incomplete batch
        
training/src/datamodules/imagenet.py