code: def main(args):
      device = f"cuda:0" if th.cuda.is_available() else "cpu"
  
      (
          g,
          labels,
          num_classes,
          split_idx,
          logger,
          train_loader,
      ) = prepare_data(args, device)
  
      embed_layer = rel_graph_embed(g, 128).to(device)
      model = EntityClassify(g, 128, num_classes).to(device)
  
      print(
          f"Number of embedding parameters: {sum(p.numel() for p in embed_layer.parameters())}"
      )
      print(
          f"Number of model parameters: {sum(p.numel() for p in model.parameters())}"
      )
  
      for run in range(args.runs):
          try:
              embed_layer.reset_parameters()
              model.reset_parameters()
          except:
              # old pytorch version doesn't support reset_parameters() API
              pass
  
          # optimizer
          all_params = itertools.chain(
              model.parameters(), embed_layer.parameters()
          )
          optimizer = th.optim.Adam(all_params, lr=0.01)
  
          if (
              args.num_workers != 0
              and device == "cpu"
              and is_support_affinity(v_t)
          ):
              expected_max = int(psutil.cpu_count(logical=False))
              if args.num_workers >= expected_max:
                  print(
                      f"[ERROR] You specified num_workers are larger than physical cores, please set any number less than {expected_max}",
                      file=sys.stderr,
                  )
              with train_loader.enable_cpu_affinity():
                  logger = train(
                      g,
                      model,
                      embed_layer,
                      optimizer,
                      train_loader,
                      split_idx,
                      labels,
                      logger,
                      device,
                      run,
                  )
          else:
              logger = train(
                  g,
                  model,
                  embed_layer,
                  optimizer,
                  train_loader,
                  split_idx,
                  labels,
                  logger,
                  device,
                  run,
              )
          logger.print_statistics(run)
  
      print("Final performance: ")
      logger.print_statistics()
  
imports: import argparse
  import itertools
  import sys
  import dgl
  import dgl.nn
  import psutil
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl import AddReverse, Compose, ToSimple
  from dgl.nn import HeteroEmbedding
  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