code: def load_pyg(name, dataset_dir):
      """
      load pyg format dataset
      :param name: dataset name
      :param dataset_dir: data directory
      :return: a list of networkx/deepsnap graphs
      """
      dataset_dir = "{}/{}".format(dataset_dir, name)
      if name in ["Cora", "CiteSeer", "PubMed"]:
          dataset_raw = Planetoid(dataset_dir, name)
      elif name[:3] == "TU_":
          # TU_IMDB doesn't have node features
          if name[3:] == "IMDB":
              name = "IMDB-MULTI"
              dataset_raw = TUDataset(
                  dataset_dir, name, transform=T.Constant()
              )
          else:
              dataset_raw = TUDataset(dataset_dir, name[3:])
          # TU_dataset only has graph-level label
          # The goal is to have synthetic tasks
          # that select smallest 100 graphs that have more than 200 edges
          if cfg.dataset.tu_simple and cfg.dataset.task != "graph":
              size = []
              for data in dataset_raw:
                  edge_num = data.edge_index.shape[1]
                  edge_num = 9999 if edge_num < 200 else edge_num
                  size.append(edge_num)
              size = torch.tensor(size)
              order = torch.argsort(size)[:100]
              dataset_raw = dataset_raw[order]
      elif name == "Karate":
          dataset_raw = KarateClub()
      elif "Coauthor" in name:
          if "CS" in name:
              dataset_raw = Coauthor(dataset_dir, name="CS")
          else:
              dataset_raw = Coauthor(dataset_dir, name="Physics")
      elif "Amazon" in name:
          if "Computers" in name:
              dataset_raw = Amazon(dataset_dir, name="Computers")
          else:
              dataset_raw = Amazon(dataset_dir, name="Photo")
      elif name == "MNIST":
          dataset_raw = MNISTSuperpixels(dataset_dir)
      elif name == "PPI":
          dataset_raw = PPI(dataset_dir)
      elif name == "QM7b":
          dataset_raw = QM7b(dataset_dir)
      else:
          raise ValueError("{} not support".format(name))
      graphs = GraphDataset.pyg_to_graphs(dataset_raw)
      return graphs
  
imports: import logging
  import pickle
  import time
  import networkx
  import torch
  import torch_geometric.transforms
  from deepsnap.batch import Batch
  from deepsnap.dataset import GraphDataset
  from ogb.graphproppred import PygGraphPropPredDataset
  from torch.utils.data import DataLoader
  from torch_geometric.datasets import PPI, Amazon, Coauthor, KarateClub, MNISTSuperpixels, Planetoid, QM7b, TUDataset
  import graphgym.models.feature_augment
  import graphgym.register
  from graphgym.config import cfg
  from graphgym.models.transform import edge_nets, ego_nets, path_len, remove_node_feature
label_desc: Use image node_id 1 for calling a global function with example usage: load_pyg(name, dataset_dir) and returns: graphs