def METHOD_NAME(self):
    if self.context.get_hparam("dataset") == "imagenet":
        valdir = os.path.join(self.download_directory, "val")
        self.normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        )
        val_dataset = datasets.ImageFolder(
            traindir,
            transforms.Compose(
                [
                    transforms.RandomResizedCrop(224),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    self.normalize,
                ]
            ),
        )
        return DataLoader(
            val_dataset,
            batch_size=self.context.get_per_slot_batch_size(),
            shuffle=False,
            num_workers=self.context.get_hparam("workers", pin_memory=True),
        )
    else:
        transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
        )
        val_dataset = datasets.CIFAR10(
            root=self.download_directory, train=False, download=True, transform=transform
        )
        return DataLoader(val_dataset, batch_size=self.context.get_per_slot_batch_size())