def setup()

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

• 12 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:
            assert self.datasets[
                "train"
            ], "In fit stage, the train dataset shouldn't be None!"
            if not self.datasets["val"] and self.val_split:
                dataset_train, dataset_val = self._split_dataset(
                    none_throws(self.datasets["train"])
                )
                # pyre-fixme[16]: `Mapping` has no attribute `__setitem__`.
                self.datasets["train"] = dataset_train
                self.datasets["val"] = dataset_val
                logging.info("We have split part of the train set into val set!")