TypeError: training_step() missing 1 required positional argument: 'dataset_idx'

[morestart]

this is my datamodule:

class MultiWindDatamodule(pl.LightningDataModule):
    def __init__(
            self,
            train_inputs_path: List[str],
            train_label_path: List[str],
            val_inputs_path: List[str],
            val_label_path: List[str],
            batch_size: int,
            num_workers: int,
            shuffle: bool):
        super(MultiWindDatamodule, self).__init__()

        self.bs = batch_size
        self.nw = num_workers
        self.shuffle = shuffle

        self.train_ds = [WindDataset2(x, y) for x, y in zip(train_inputs_path, train_label_path)]
        self.val_ds = [WindDataset2(x, y) for x, y in zip(val_inputs_path, val_label_path)]
        print()

    def train_dataloader(self) -> TRAIN_DATALOADERS:
        return [DataLoader(
            ds,
            batch_size=self.bs,
            num_workers=self.nw,
            shuffle=self.shuffle
        ) for ds in self.train_ds]

    def val_dataloader(self) -> EVAL_DATALOADERS:
        return [DataLoader(
            ds,
            batch_size=self.bs,
            num_workers=self.nw,
            shuffle=self.shuffle
        ) for ds in self.val_ds]

train step and val step:

def training_step(self, batch, batch_idx, dataset_idx):
    x, y = batch
    pre = self.forward(x)

    loss = self.loss(pre, y)
    self.log('train_loss', loss)

    return loss

def validation_step(self, batch, batch_idx, dataloader_idx):
    x, y = batch
    pre = self.forward(x)

    loss = self.loss(pre, y)

    self.log('val_loss', loss)

    return loss

then get this error:
TypeError: training_step() missing 1 required positional argument: 'dataset_idx'

[morestart]

anyone know this?

[akihironitta]

@morestart I believe this patch will make it work:

-def training_step(self, batch, batch_idx, dataset_idx):
+def training_step(self, batch, batch_idx, dataloader_idx):