Method __init__

tests/test_pytorch_lightning.py:13–16

(self, batch_size=16)

class LitDataModule(pl.LightningDataModule):

    def __init__(self, batch_size=16):
        super().__init__()

        self.batch_size = batch_size

    def setup(self, stage=None):
        X_train = torch.rand(100, 1, 28, 28)