The following is a working Pytorch Lightning DataModule.
import os from pytorch_lightning import LightningDataModule import torchvision.datasets as datasets f