Method setup

data/landoceandataset.py:84–87

(self, stage: str)

        self.batch_size=batch_size

    def setup(self, stage: str):
        self.train_ds = TensorDataset(*get_data(self.num_samples, seed=0))
        self.valid_ds = TensorDataset(*get_data(self.num_samples, seed=1))
        self.evalu_ds = TensorDataset(*get_data(self.num_samples, grid=True))

    def train_dataloader(self):
        return DataLoader(self.train_ds, batch_size=self.batch_size, shuffle=True)