code: def main(args):
      device = (
          "cuda:0"
          if torch.cuda.is_available() and args.num_gpus > 0
          else "cpu"
      )
  
      # Prepare the data.
      (
          g,
          labels,
          num_classes,
          split_idx,
          train_loader,
          feats,
      ) = prepare_data(args, device)
  
      feat_size = 128 if args.dataset == "ogbn-mag" else 768
  
      # Create the embedding layer and move it to the appropriate device.
      embed_layer = None
      if args.dataset == "ogbn-mag":
          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())}"
      )
  
      try:
          if embed_layer is not None:
              embed_layer.reset_parameters()
          model.reset_parameters()
      except:
          # Old pytorch version doesn't support reset_parameters() API.
          ##################################################################
          # [Why we need to reset the parameters?]
          # If parameters are not reset, the model will start with the
          # parameters learned from the last run, potentially resulting
          # in biased outcomes or sub-optimal performance if the model was
          # previously stuck in a poor local minimum.
          ##################################################################
          pass
  
      # `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`` is the number of physical cores on your machine.
      # The `logical` parameter, when set to False, ensures that the count
      # returned is the number of physical cores instead of logical cores
      # (which could be higher due to technologies like Hyper-Threading).
      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,
          feats,
          model,
          embed_layer,
          optimizer,
          train_loader,
          split_idx,
          labels,
          device,
      )
  
      print("Testing...")
      test_key = "test" if args.dataset == "ogbn-mag" else "test-dev"
      test_acc = evaluate(
          args.dataset,
          g,
          feats,
          model,
          embed_layer,
          labels,
          device,
          split_idx[test_key],
          save_test_submission=(args.dataset == "ogb-lsc-mag240m"),
      )
      print(f"Test accuracy {test_acc*100:.4f}")
  
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 7 for calling a global function with example usage: main(args) without return types