class_docstr: Set up Pytorch DataSet and DataLoader objects to prepare data for single-modal/multimodal training,
  validation, testing, and prediction. We organize the multimodal data using pd.DataFrame.
  For some modalities, e.g, image, that cost much memory, we only store their disk path to do lazy loading.
  This class inherits from the Pytorch Lightning's LightningDataModule:
  https://lightning.ai/docs/pytorch/stable/data/datamodule.html
class_docstr_tok: ["Set","up","Pytorch","DataSet","and","DataLoader","objects","to","prepare","data","for","single-modal\/multimodal","training",",","validation",",","testing",",","and","prediction",".","We","organize","the","multimodal","data","using","pd.DataFrame",".","For","some","modalities",",","e.g",",","image",",","that","cost","much","memory",",","we","only","store","their","disk","path","to","do","lazy","loading",".","This","class","inherits","from","the","Pytorch","Lightning","'s","LightningDataModule",":","https",":","\/\/lightning.ai\/docs\/pytorch\/stable\/data\/datamodule.html"]
docstr: Set up datasets for different stages: "fit" (training and validation), "test", and "predict".
  This method is registered by Pytorch Lightning's LightningDataModule.
  Refer to: https://lightning.ai/docs/pytorch/stable/data/datamodule.html#setup
  
  Parameters
  ----------
  stage
      Stage name including choices:
          - fit (For the fitting stage)
          - test (For the test stage)
          - predict (For the prediction stage)
docstr_tok: ["Set","up","datasets","for","different","stages",":","``","fit","''","(","training","and","validation",")",",","``","test","''",",","and","``","predict","''",".","This","method","is","registered","by","Pytorch","Lightning","'s","LightningDataModule",".","Refer","to",":","https",":","\/\/lightning.ai\/docs\/pytorch\/stable\/data\/datamodule.html","#","setup","Parameters","--","--","--","--","--","stage","Stage","name","including","choices",":","-","fit","(","For","the","fitting","stage",")","-","test","(","For","the","test","stage",")","-","predict","(","For","the","prediction","stage",")"]
code: def setup(self, stage):
      """
      Set up datasets for different stages: "fit" (training and validation), "test", and "predict".
      This method is registered by Pytorch Lightning's LightningDataModule.
      Refer to: https://lightning.ai/docs/pytorch/stable/data/datamodule.html#setup
  
      Parameters
      ----------
      stage
          Stage name including choices:
              - fit (For the fitting stage)
              - test (For the test stage)
              - predict (For the prediction stage)
      """
      if stage == "fit":
          self.set_dataset(TRAIN)
          self.set_dataset(VALIDATE)
      elif stage == "validate":
          self.set_dataset(VALIDATE)
      elif stage == "test":
          self.set_dataset(TEST)
      elif stage == "predict":
          self.set_dataset(PREDICT)
      else:
          raise ValueError(f"Unknown stage {stage}")
  
file_path: autogluon/multimodal/src/autogluon/multimodal/data/datamodule.py
imports: from typing import Dict, List, Optional, Union
  import pandas
  from lightning.pytorch import LightningDataModule
  from torch.utils.data import DataLoader, Dataset
  from ..constants import PREDICT, TEST, TRAIN, VALIDATE
  from .dataset import BaseDataset
  from .preprocess_dataframe import MultiModalFeaturePreprocessor
  from .utils import get_collate_fn
label_desc: Use image node_id 3 for calling the BaseDataModule obj's underlying member method code with example usage: obj.setup(stage) without return types