code: def prepare_data(args, device):
      feats = {}
      if args.dataset == "ogbn-mag":
          dataset = DglNodePropPredDataset(
              name="ogbn-mag", root=args.rootdir
          )
  
          # - graph: dgl graph object.
          # - label: torch tensor of shape (num_nodes, num_tasks).
          g, labels = dataset[0]
  
          # Flatten the labels for "paper" type nodes. This step reduces the
          # dimensionality of the labels. We need to flatten the labels because
          # the model requires a 1-dimensional label tensor.
          labels = labels["paper"].flatten().long()
  
          # Apply transformation to the graph.
          # - "ToSimple()" removes multi-edge between two nodes.
          # - "AddReverse()" adds reverse edges to the graph.
          print("Start to transform graph. This may take a while...")
          transform = Compose([ToSimple(), AddReverse()])
          g = transform(g)
      else:
          dataset = MAG240MDataset(root=args.rootdir)
          (g,), _ = dgl.load_graphs(args.graph_path)
          g = g.formats(["csc"])
          labels = torch.as_tensor(dataset.paper_label).long()
          # As feature data is too large to fit in memory, we read it from disk.
          feats["paper"] = torch.as_tensor(
              np.load(args.paper_feature_path, mmap_mode="r+")
          )
          feats["author"] = torch.as_tensor(
              np.load(args.author_feature_path, mmap_mode="r+")
          )
          feats["institution"] = torch.as_tensor(
              np.load(args.inst_feature_path, mmap_mode="r+")
          )
      print(f"Loaded graph: {g}")
  
      # Get train/valid/test index.
      split_idx = dataset.get_idx_split()
      if args.dataset == "ogb-lsc-mag240m":
          split_idx = {
              split_type: {"paper": split_idx[split_type]}
              for split_type in split_idx
          }
  
      # Initialize a train sampler that samples neighbors for multi-layer graph
      # convolution. It samples 25 and 10 neighbors for the first and second
      # layers respectively.
      sampler = dgl.dataloading.MultiLayerNeighborSampler([25, 10])
      num_workers = args.num_workers
      train_loader = dgl.dataloading.DataLoader(
          g,
          split_idx["train"],
          sampler,
          batch_size=1024,
          shuffle=True,
          num_workers=num_workers,
          device=device,
      )
  
      return (
          g,
          labels,
          dataset.num_classes,
          split_idx,
          train_loader,
          feats,
      )
  
imports: import argparse
  import itertools
  import sys
  import time
  import dgl
  import dgl.nn
  import numpy
  import psutil
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl import AddReverse, Compose, ToSimple
  from dgl.nn import HeteroEmbedding
  from ogb.lsc import MAG240MDataset, MAG240MEvaluator
  from ogb.nodeproppred import DglNodePropPredDataset, Evaluator
  from tqdm import tqdm
label_desc: Use image node_id 1 for calling a global function with example usage: prepare_data(args, device) and returns: g,  labels,  dataset,  split_idx,  train_loader,  feats