code: def load_ogb(dataset):
      if dataset == "ogbn-mag":
          dataset = DglNodePropPredDataset(name=dataset)
          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"
          ]
          paper_labels = labels["paper"].squeeze()
  
          num_rels = len(hg.canonical_etypes)
          num_of_ntype = len(hg.ntypes)
          num_classes = dataset.num_classes
          category = "paper"
          print("Number of relations: {}".format(num_rels))
          print("Number of class: {}".format(num_classes))
          print("Number of train: {}".format(len(train_idx)))
          print("Number of valid: {}".format(len(val_idx)))
          print("Number of test: {}".format(len(test_idx)))
  
          # get target category id
          category_id = len(hg.ntypes)
          for i, ntype in enumerate(hg.ntypes):
              if ntype == category:
                  category_id = i
  
          train_mask = th.zeros((hg.num_nodes("paper"),), dtype=th.bool)
          train_mask[train_idx] = True
          val_mask = th.zeros((hg.num_nodes("paper"),), dtype=th.bool)
          val_mask[val_idx] = True
          test_mask = th.zeros((hg.num_nodes("paper"),), dtype=th.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
  
          hg.nodes["paper"].data["labels"] = paper_labels
          return hg
      else:
          raise ("Do not support other ogbn datasets.")
  
imports: import argparse
  import time
  import dgl
  import numpy
  import torch
  from ogb.nodeproppred import DglNodePropPredDataset
label_desc: Use image node_id 1 for calling a global function with example usage: load_ogb(dataset) and returns: hg