def setup(self, stage: Optional[str] = None) -> None:
        """Creates train, val, and test dataset."""
        if stage == "fit" or stage is None:
            train_transforms = (self.train_transform() if self.train_transforms is None
                                else self.train_transforms)
            val_transforms = (self.val_transform() if self.val_transforms is None
                              else self.val_transforms)
            if self.img_dtype is not torch.float32:
                assert isinstance(train_transforms, transforms.Compose)
                assert isinstance(val_transforms, transforms.Compose)
                convert_dtype = transforms.Lambda(lambda x: x.to(dtype=self.img_dtype))
                train_transforms.transforms.append(convert_dtype)
                val_transforms.transforms.append(convert_dtype)
            self.dataset_train = ImageFolder(self.data_dir / 'train', transform=train_transforms)
            self.dataset_val = ImageFolder(self.data_dir / 'val', transform=val_transforms)

        if stage == "test" or stage is None:
            test_transforms = (self.val_transform() if self.test_transforms is None
                               else self.test_transforms)
            if self.img_dtype is not torch.float32:
                assert isinstance(test_transforms, transforms.Compose)
                convert_dtype = transforms.Lambda(lambda x: x.to(dtype=self.img_dtype))
                test_transforms.transforms.append(convert_dtype)
            self.dataset_test = ImageFolder(self.data_dir / 'val', transform=test_transforms)
Creates train, val, and test dataset.
training/src/datamodules/imagenet.py