code: def load_ogb_mag():
      name = "ogbn-mag"
      os.symlink("/tmp/dataset/", os.path.join(os.getcwd(), "dataset"))
  
      print("load", name)
      dataset = DglNodePropPredDataset(name=name)
      print("finish loading", name)
      split_idx = dataset.get_idx_split()
      train_idx = split_idx["train"]["paper"]
      val_idx = split_idx["valid"]["paper"]
      test_idx = split_idx["test"]["paper"]
      hg_orig, labels = dataset[0]
      subgs = {}
      for etype in hg_orig.canonical_etypes:
          u, v = hg_orig.all_edges(etype=etype)
          subgs[etype] = (u, v)
          subgs[(etype[2], "rev-" + etype[1], etype[0])] = (v, u)
      hg = dgl.heterograph(subgs)
      hg.nodes["paper"].data["feat"] = hg_orig.nodes["paper"].data[
          "feat"
      ]
      hg.nodes["paper"].data["labels"] = labels["paper"].squeeze()
      train_mask = torch.zeros(
          (hg.num_nodes("paper"),), dtype=torch.bool
      )
      train_mask[train_idx] = True
      val_mask = torch.zeros((hg.num_nodes("paper"),), dtype=torch.bool)
      val_mask[val_idx] = True
      test_mask = torch.zeros(
          (hg.num_nodes("paper"),), dtype=torch.bool
      )
      test_mask[test_idx] = True
      hg.nodes["paper"].data["train_mask"] = train_mask
      hg.nodes["paper"].data["val_mask"] = val_mask
      hg.nodes["paper"].data["test_mask"] = test_mask
  
      num_classes = dataset.num_classes
      return OGBDataset(hg, num_classes, "paper")
  
imports: import inspect
  import json
  import os
  import pickle
  import shutil
  import time
  import zipfile
  from functools import partial, reduce, wraps
  from timeit import default_timer
  import dgl
  import numpy
  import pandas
  import requests
  import torch
  from ogb.nodeproppred import DglNodePropPredDataset
  import traceback
  from _thread import start_new_thread
  import torch.multiprocessing
label_desc: Use image node_id 8 for calling a global function with example usage: load_ogb_mag() and returns: OGBDataset