code: def __init__(
      self,
      data_dir: str = "./",
      batch_size: int = 32,
      use_trials: bool = False,
  ) -> None:
      super().__init__()
  
      self.data_dir = data_dir
      self.batch_size = batch_size
  
      # TrialMNIST is a constrained MNIST dataset
      self.dataset_cls = TrialMNIST if use_trials else MNIST
  
imports: import torch
  from lightning.pytorch.core.datamodule import LightningDataModule
  from lightning_utilities.core.imports import RequirementCache
  from torch.utils.data import DataLoader
  from tests_pytorch.helpers.datasets import MNIST, SklearnDataset, TrialMNIST
label_desc: Use image node_id 1 to create a new MNISTDataModule object from inherited base classes: LightningDataModule with example: obj = MNISTDataModule(data_dir, batch_size, use_trials)