def val_dataloader(self):
        if self.hparams.dataset_name == 'CIFAR10':
            dataset = AMDIMPretraining.cifar10_val(self.hparams.data_dir)

        if self.hparams.dataset_name == 'stl_10':
            dataset = self.val_split

        if self.hparams.dataset_name == 'imagenet_128':
            dataset = AMDIMPretraining.imagenet_val(self.hparams.data_dir, self.hparams.nb_classes)

        # LOADER
        loader = DataLoader(
            dataset=dataset,
            batch_size=self.hparams.batch_size,
            pin_memory=True,
            drop_last=True,
            num_workers=16,
        )
        return loader
    def val_dataloader(self):
        kwargs = dict(nb_classes=self.hparams.nb_classes) if self.hparams.datamodule == 'imagenet2012' else {}
        dataset = AMDIMPretraining.get_dataset(self.hparams.datamodule, self.hparams.data_dir, split='val', **kwargs)

        # LOADER
        loader = DataLoader(
            dataset=dataset,
            batch_size=self.hparams.batch_size,
            pin_memory=True,
            drop_last=True,
            num_workers=16,
        )
        return loader