def setup(self, stage=None):
        """Creates train, val, and test dataset."""

        from typing import Any, Callable, List, Optional, Union

        import hydra  # for mixup
        from pl_bolts.transforms.dataset_normalizations import \
            imagenet_normalization
        from torch.utils.data import Dataset
        from torch.utils.data.dataloader import default_collate
        from torchvision.datasets import ImageFolder

        # for access in other methods
        self.imagenet_normalization = imagenet_normalization
        self.default_collate = default_collate
        self.hydra = hydra
        self.ImageFolder = ImageFolder

        if self.mixup is not None:
            self.mixup_fn = hydra.utils.instantiate(self.mixup)
        else:
            self.mixup_fn = None

        self.dir_path = self.data_dir or default_data_path / self._name_

        if stage == "fit" or stage is None:
            self.set_phase([self.image_size])

        #     train_transforms = (self.train_transform() if self.train_transforms is None
        #                         else hydra.utils.instantiate(self.train_transforms))
        #     val_transforms = (self.val_transform() if self.val_transforms is None
        #                       else hydra.utils.instantiate(self.val_transforms))

        #     self.dataset_train = ImageFolder(self.dir_path / 'val',  # modded
        #                                         transform=train_transforms)

        #     if self.val_split > 0.:
        #         # this will create the val split
        #         self.split_train_val(self.val_split)
        #     # will use the test split as val by default
        #     else:
        #         self.dataset_val = ImageFolder(self.dir_path / 'val', transform=val_transforms)

        #     # modded, override (for debugging)
        #     self.dataset_train = self.dataset_val

        if stage == "test" or stage is None:
            test_transforms = (self.val_transform() if self.test_transforms is None
                               else hydra.utils.instantiate(self.test_transforms))

            self.dataset_test = ImageFolder(os.path.join(self.dir_path, 'val'), transform=test_transforms)

            # # modded, override (for debugging)
            # self.dataset_test = self.dataset_val
Creates train, val, and test dataset.
analyze/convnexts4nd/src/dataloaders/vision.py