code: def setup(self, stage=None):
      # Assign train/val datasets for use in dataloaders
      if stage == "fit" or stage is None:
          iris_full = self._get_iris_as_tensor_dataset()
          self.train_set, self.val_set = random_split(
              iris_full, [130, 20]
          )
  
      # Assign test dataset for use in dataloader(s)
      if stage == "test" or stage is None:
          self.train_set, self.test_set = random_split(
              self.train_set, [110, 20]
          )
  
file_path: mlflow/tests/pytorch/iris_data_module.py
imports: import pytorch_lightning
  import torch
  from sklearn.datasets import load_iris
  from torch.utils.data import DataLoader, TensorDataset, random_split
label_desc: Use image node_id 3 for calling the IrisDataModuleBase obj's underlying member method code with example usage: obj.setup(stage) without return types