2860e37da4e0b1335f25fcb0f62ab2d81698c752,examples/_tests_mnist_stages2/experiment.py,Experiment,get_datasets,#Experiment#Any#,19

Before Change

            download=True,
            transform=Experiment.get_transforms(mode="train", stage=stage)
        )
        testset = torchvision.datasets.MNIST(
            "./data",
            train=False,
            download=True,
            transform=Experiment.get_transforms(mode="train", stage=stage)
        )

        datasets["train"] = trainset
        datasets["valid"] = testset

After Change

class Experiment(ConfigExperiment):
    def get_datasets(self, stage: str, **kwargs):
        datasets = OrderedDict()
        for mode in ("train", "valid"):
            datasets[mode] = MNIST(
                "./data",
                train=False,
                download=True,
                transform=self.get_transforms(stage=stage, dataset=mode)
            )

        return datasets

Instances