code: def set_fold(
      dataset: InMemoryDataset, fold_id: int
  ) -> InMemoryDataset:
      r"""Set fold for graph dataset consist of multiple graphs.
  
      Parameters
      ----------
      dataset: `autogl.data.InMemoryDataset`
          dataset with multiple graphs.
      fold_id: `int`
          The fold in to use, MUST be in [0, dataset.n_splits)
  
      Returns
      -------
      `autogl.data.InMemoryDataset`
          The reference of original dataset.
      """
      if not (hasattr(dataset, "folds") and dataset.folds is not None):
          raise ValueError("Dataset do NOT contain folds")
      if not 0 <= fold_id < len(dataset.folds):
          raise ValueError(
              f"Fold id {fold_id} exceed total cross validation split number {len(dataset.folds)}"
          )
      dataset.train_index = dataset.folds[fold_id][0]
      dataset.val_index = dataset.folds[fold_id][1]
      setattr(
          dataset,
          "cross_train_split",
          [dataset[i] for i in dataset.train_index],
      )
      setattr(
          dataset,
          "cross_val_split",
          [dataset[i] for i in dataset.val_index],
      )
      return dataset
  
file_path: AutoGL/autogl/datasets/utils/_general.py
imports: import numpy
  import random
  import torch
  import torch.utils.data
  import typing
  from sklearn.model_selection import StratifiedKFold, KFold
  from autogl import backend
  from autogl.data import InMemoryDataset
label_desc: Use image node_id 5 for calling a global function with example usage: set_fold(dataset, fold_id) and returns: dataset