code: def main(args):
      device = (
          torch.device("cuda")
          if args.num_gpus > 0
          else torch.device("cpu")
      )
  
      # Load dataset.
      (
          g,
          features,
          train_set,
          valid_set,
          test_set,
          num_classes,
      ) = load_dataset(args.dataset)
  
      feat_size = features.size("node", "paper", "feat")[0]
  
      # As `ogb-lsc-mag240m` is a large dataset, features of `author` and
      # `institution` are generated in advance and stored in the feature store.
      # For `ogbn-mag`, we generate the features on the fly.
      embed_layer = None
      if args.dataset == "ogbn-mag":
          # Create the embedding layer and move it to the appropriate device.
          embed_layer = rel_graph_embed(g, feat_size).to(device)
          print(
              "Number of embedding parameters: "
              f"{sum(p.numel() for p in embed_layer.parameters())}"
          )
  
      # Initialize the entity classification model.
      model = EntityClassify(g, feat_size, num_classes).to(device)
  
      print(
          "Number of model parameters: "
          f"{sum(p.numel() for p in model.parameters())}"
      )
  
      if embed_layer is not None:
          embed_layer.reset_parameters()
      model.reset_parameters()
  
      # `itertools.chain()` is a function in Python's itertools module.
      # It is used to flatten a list of iterables, making them act as
      # one big iterable.
      # In this context, the following code is used to create a single
      # iterable over the parameters of both the model and the embed_layer,
      # which is passed to the optimizer. The optimizer then updates all
      # these parameters during the training process.
      all_params = itertools.chain(
          model.parameters(),
          [] if embed_layer is None else embed_layer.parameters(),
      )
      optimizer = torch.optim.Adam(all_params, lr=0.01)
  
      expected_max = int(psutil.cpu_count(logical=False))
      if args.num_workers >= expected_max:
          print(
              "[ERROR] You specified num_workers are larger than physical"
              f"cores, please set any number less than {expected_max}",
              file=sys.stderr,
          )
  
      train(
          args.dataset,
          g,
          model,
          embed_layer,
          optimizer,
          train_set,
          valid_set,
          device,
          features,
          args.num_workers,
          args.num_epochs,
      )
  
      print("Testing...")
      test_acc = evaluate(
          args.dataset,
          g,
          model,
          embed_layer,
          device,
          test_set,
          features,
          args.num_workers,
      )
      print(f"Test accuracy {test_acc*100:.4f}")
  
imports: import argparse
  import itertools
  import sys
  import time
  import dgl
  import dgl.graphbolt
  import dgl.nn
  import psutil
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl.nn import HeteroEmbedding
  from ogb.lsc import MAG240MEvaluator
  from ogb.nodeproppred import Evaluator
  from tqdm import tqdm
label_desc: Use image node_id 8 for calling a global function with example usage: main(args) without return types