def setup()

in torchrecipes/audio/source_separation/datamodule/librimix.py [0:0]

• 27 lines of code
• 8 McCabe index (conditional complexity)

    def setup(self, stage: Optional[str] = None):

        # Assign train/val datasets for use in dataloaders
        if stage == "fit" or stage is None:
            if self.testing:
                self.train = TestDataset()
                self.val = TestDataset()
            else:
                self.train = LibriMix(
                    self.root_dir,
                    self.tr_split,
                    self.num_speakers,
                    self.sample_rate,
                    self.task,
                )
                self.val = LibriMix(
                    self.root_dir, "dev", self.num_speakers, self.sample_rate, self.task
                )

        if stage == "test" or stage is None:
            if self.testing:
                self.test = TestDataset()
            else:
                self.test = LibriMix(
                    self.root_dir,
                    "test",
                    self.num_speakers,
                    self.sample_rate,
                    self.task,
                )