def setup()

in torchrecipes/vision/data/modules/mnist_data_module.py [0:0]

• 14 lines of code
• 5 McCabe index (conditional complexity)

    def setup(self, stage: Optional[str] = None) -> None:
        """Creates train, val and test dataset."""
        if stage == "fit" or stage is None:
            dataset_train = MNIST(
                self.data_dir, train=True, transform=self.train_transforms
            )
            dataset_val = MNIST(
                self.data_dir, train=True, transform=self.val_transforms
            )

            # Split
            self.datasets["train"] = self._split_dataset(dataset_train)
            self.datasets["val"] = self._split_dataset(dataset_val, train=False)

        if stage == "test" or stage is None:
            self.datasets["test"] = MNIST(
                self.data_dir, train=False, transform=self.test_transforms
            )