code: def load_pyg(name, dataset_dir):
      """Load PyG dataset objects. (More PyG datasets will be supported).
  
      Args:
          name (str): dataset name
          dataset_dir (str): data directory
  
      Returns: PyG dataset object
  
      """
      dataset_dir = osp.join(dataset_dir, name)
      if name in ["Cora", "CiteSeer", "PubMed"]:
          dataset = Planetoid(dataset_dir, name)
      elif name[:3] == "TU_":
          # TU_IMDB doesn't have node features
          if name[3:] == "IMDB":
              name = "IMDB-MULTI"
              dataset = TUDataset(
                  dataset_dir, name, transform=T.Constant()
              )
          else:
              dataset = TUDataset(dataset_dir, name[3:])
      elif name == "Karate":
          dataset = KarateClub()
      elif "Coauthor" in name:
          if "CS" in name:
              dataset = Coauthor(dataset_dir, name="CS")
          else:
              dataset = Coauthor(dataset_dir, name="Physics")
      elif "Amazon" in name:
          if "Computers" in name:
              dataset = Amazon(dataset_dir, name="Computers")
          else:
              dataset = Amazon(dataset_dir, name="Photo")
      elif name == "MNIST":
          dataset = MNISTSuperpixels(dataset_dir)
      elif name == "PPI":
          dataset = PPI(dataset_dir)
      elif name == "QM7b":
          dataset = QM7b(dataset_dir)
      else:
          raise ValueError(f"'{name}' not support")
  
      return dataset
  
file_path: pytorch_geometric/torch_geometric/graphgym/loader.py
imports: import os.path
  from typing import Callable
  import torch
  import torch_geometric.graphgym.register
  import torch_geometric.transforms
  from torch_geometric.datasets import PPI, Amazon, Coauthor, KarateClub, MNISTSuperpixels, Planetoid, QM7b, TUDataset
  from torch_geometric.graphgym.config import cfg
  from torch_geometric.graphgym.models.transform import create_link_label, neg_sampling_transform
  from torch_geometric.loader import ClusterLoader, DataLoader, GraphSAINTEdgeSampler, GraphSAINTNodeSampler, GraphSAINTRandomWalkSampler, NeighborSampler, RandomNodeLoader
  from torch_geometric.utils import index_to_mask, negative_sampling, to_undirected
label_desc: Use image node_id 2 for calling a global function with example usage: load_pyg(name, dataset_dir) and returns: dataset