code: def setup(self, stage: str):
      if stage == "fit":
          self.mnist_train = self.dataset_cls(self.data_dir, train=True)
      if stage == "test":
          self.mnist_test = self.dataset_cls(self.data_dir, train=False)
  
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 3 for calling the MNISTDataModule obj's underlying member method code with example usage: obj.setup(stage) without return types