code: def __init__(self, batch_size: int = 32):
      super().__init__()
      dataset = MNIST(
          DATASETS_PATH,
          train=True,
          download=True,
          transform=transforms.ToTensor(),
      )
      self.mnist_test = MNIST(
          DATASETS_PATH,
          train=False,
          download=True,
          transform=transforms.ToTensor(),
      )
      self.mnist_train, self.mnist_val = random_split(
          dataset,
          [55000, 5000],
          generator=torch.Generator().manual_seed(42),
      )
      self.batch_size = batch_size
  
imports: from os import path
  from typing import Optional, Tuple
  import torch
  import torch.nn.functional
  from lightning.pytorch import LightningDataModule, LightningModule, Trainer, callbacks, cli_lightning_logo
  from lightning.pytorch.cli import LightningCLI
  from lightning.pytorch.demos.mnist_datamodule import MNIST
  from lightning.pytorch.utilities import rank_zero_only
  from lightning.pytorch.utilities.imports import _TORCHVISION_AVAILABLE
  from torch import nn
  from torch.utils.data import DataLoader, random_split
label_desc: Use image node_id 1 to create a new MyDataModule object from inherited base classes: LightningDataModule with example: obj = MyDataModule(batch_size)