code: def run(args, device, data):
      (
          g,
          num_classes,
          train_nid,
          val_nid,
          test_nid,
          labels,
          all_val_nid,
          all_test_nid,
      ) = data
  
      fanouts = [int(fanout) for fanout in args.fanout.split(",")]
      val_fanouts = [
          int(fanout) for fanout in args.validation_fanout.split(",")
      ]
  
      sampler = dgl.dataloading.MultiLayerNeighborSampler(fanouts)
      dataloader = dgl.dataloading.DistNodeDataLoader(
          g,
          {"paper": train_nid},
          sampler,
          batch_size=args.batch_size,
          shuffle=True,
          drop_last=False,
      )
  
      valid_sampler = dgl.dataloading.MultiLayerNeighborSampler(
          val_fanouts
      )
      valid_dataloader = dgl.dataloading.DistNodeDataLoader(
          g,
          {"paper": val_nid},
          valid_sampler,
          batch_size=args.batch_size,
          shuffle=False,
          drop_last=False,
      )
  
      test_sampler = dgl.dataloading.MultiLayerNeighborSampler(
          val_fanouts
      )
      test_dataloader = dgl.dataloading.DistNodeDataLoader(
          g,
          {"paper": test_nid},
          test_sampler,
          batch_size=args.eval_batch_size,
          shuffle=False,
          drop_last=False,
      )
  
      embed_layer = DistEmbedLayer(
          device,
          g,
          args.n_hidden,
          sparse_emb=args.sparse_embedding,
          dgl_sparse_emb=args.dgl_sparse,
          feat_name="feat",
      )
  
      model = EntityClassify(
          device,
          args.n_hidden,
          num_classes,
          g.etypes,
          num_bases=args.n_bases,
          num_hidden_layers=args.n_layers - 2,
          dropout=args.dropout,
          use_self_loop=args.use_self_loop,
          layer_norm=args.layer_norm,
      )
      model = model.to(device)
  
      if not args.standalone:
          if args.num_gpus == -1:
              model = DistributedDataParallel(model)
              # If there are dense parameters in the embedding layer
              # or we use Pytorch saprse embeddings.
              if len(embed_layer.node_projs) > 0 or not args.dgl_sparse:
                  embed_layer = DistributedDataParallel(embed_layer)
          else:
              dev_id = g.rank() % args.num_gpus
              model = DistributedDataParallel(
                  model, device_ids=[dev_id], output_device=dev_id
              )
              # If there are dense parameters in the embedding layer
              # or we use Pytorch saprse embeddings.
              if len(embed_layer.node_projs) > 0 or not args.dgl_sparse:
                  embed_layer = embed_layer.to(device)
                  embed_layer = DistributedDataParallel(
                      embed_layer,
                      device_ids=[dev_id],
                      output_device=dev_id,
                  )
  
      if args.sparse_embedding:
          if args.dgl_sparse and args.standalone:
              emb_optimizer = dgl.distributed.optim.SparseAdam(
                  list(embed_layer.node_embeds.values()),
                  lr=args.sparse_lr,
              )
              print(
                  "optimize DGL sparse embedding:",
                  embed_layer.node_embeds.keys(),
              )
          elif args.dgl_sparse:
              emb_optimizer = dgl.distributed.optim.SparseAdam(
                  list(embed_layer.module.node_embeds.values()),
                  lr=args.sparse_lr,
              )
              print(
                  "optimize DGL sparse embedding:",
                  embed_layer.module.node_embeds.keys(),
              )
          elif args.standalone:
              emb_optimizer = th.optim.SparseAdam(
                  list(embed_layer.node_embeds.parameters()),
                  lr=args.sparse_lr,
              )
              print(
                  "optimize Pytorch sparse embedding:",
                  embed_layer.node_embeds,
              )
          else:
              emb_optimizer = th.optim.SparseAdam(
                  list(embed_layer.module.node_embeds.parameters()),
                  lr=args.sparse_lr,
              )
              print(
                  "optimize Pytorch sparse embedding:",
                  embed_layer.module.node_embeds,
              )
  
          dense_params = list(model.parameters())
          if args.standalone:
              dense_params += list(embed_layer.node_projs.parameters())
              print(
                  "optimize dense projection:", embed_layer.node_projs
              )
          else:
              dense_params += list(
                  embed_layer.module.node_projs.parameters()
              )
              print(
                  "optimize dense projection:",
                  embed_layer.module.node_projs,
              )
          optimizer = th.optim.Adam(
              dense_params, lr=args.lr, weight_decay=args.l2norm
          )
      else:
          all_params = list(model.parameters()) + list(
              embed_layer.parameters()
          )
          optimizer = th.optim.Adam(
              all_params, lr=args.lr, weight_decay=args.l2norm
          )
  
      # training loop
      print("start training...")
      for epoch in range(args.n_epochs):
          tic = time.time()
  
          sample_time = 0
          copy_time = 0
          forward_time = 0
          backward_time = 0
          update_time = 0
          number_train = 0
          number_input = 0
  
          step_time = []
          iter_t = []
          sample_t = []
          feat_copy_t = []
          forward_t = []
          backward_t = []
          update_t = []
          iter_tput = []
  
          start = time.time()
          # Loop over the dataloader to sample the computation dependency graph as a list of
          # blocks.
          step_time = []
          for step, sample_data in enumerate(dataloader):
              input_nodes, seeds, blocks = sample_data
              seeds = seeds["paper"]
              number_train += seeds.shape[0]
              number_input += np.sum(
                  [
                      blocks[0].num_src_nodes(ntype)
                      for ntype in blocks[0].ntypes
                  ]
              )
              tic_step = time.time()
              sample_time += tic_step - start
              sample_t.append(tic_step - start)
  
              feats = embed_layer(input_nodes)
              label = labels[seeds].to(device)
              copy_time = time.time()
              feat_copy_t.append(copy_time - tic_step)
  
              # forward
              logits = model(blocks, feats)
              assert len(logits) == 1
              logits = logits["paper"]
              loss = F.cross_entropy(logits, label)
              forward_end = time.time()
  
              # backward
              optimizer.zero_grad()
              if args.sparse_embedding:
                  emb_optimizer.zero_grad()
              loss.backward()
              compute_end = time.time()
              forward_t.append(forward_end - copy_time)
              backward_t.append(compute_end - forward_end)
  
              # Update model parameters
              optimizer.step()
              if args.sparse_embedding:
                  emb_optimizer.step()
              update_t.append(time.time() - compute_end)
              step_t = time.time() - start
              step_time.append(step_t)
  
              train_acc = th.sum(
                  logits.argmax(dim=1) == label
              ).item() / len(seeds)
  
              if step % args.log_every == 0:
                  print(
                      "[{}] Epoch {:05d} | Step {:05d} | Train acc {:.4f} | Loss {:.4f} | time {:.3f} s"
                      "| sample {:.3f} | copy {:.3f} | forward {:.3f} | backward {:.3f} | update {:.3f}".format(
                          g.rank(),
                          epoch,
                          step,
                          train_acc,
                          loss.item(),
                          np.sum(step_time[-args.log_every :]),
                          np.sum(sample_t[-args.log_every :]),
                          np.sum(feat_copy_t[-args.log_every :]),
                          np.sum(forward_t[-args.log_every :]),
                          np.sum(backward_t[-args.log_every :]),
                          np.sum(update_t[-args.log_every :]),
                      )
                  )
              start = time.time()
  
          gc.collect()
          print(
              "[{}]Epoch Time(s): {:.4f}, sample: {:.4f}, data copy: {:.4f}, forward: {:.4f}, backward: {:.4f}, update: {:.4f}, #train: {}, #input: {}".format(
                  g.rank(),
                  np.sum(step_time),
                  np.sum(sample_t),
                  np.sum(feat_copy_t),
                  np.sum(forward_t),
                  np.sum(backward_t),
                  np.sum(update_t),
                  number_train,
                  number_input,
              )
          )
          epoch += 1
  
          start = time.time()
          g.barrier()
          val_acc, test_acc = evaluate(
              g,
              model,
              embed_layer,
              labels,
              valid_dataloader,
              test_dataloader,
              all_val_nid,
              all_test_nid,
          )
          if val_acc >= 0:
              print(
                  "Val Acc {:.4f}, Test Acc {:.4f}, time: {:.4f}".format(
                      val_acc, test_acc, time.time() - start
                  )
              )
  
imports: import argparse
  import gc, os
  import itertools
  import time
  import numpy
  from functools import partial
  import dgl
  import torch
  import torch.multiprocessing
  import torch.nn
  import torch.nn.functional
  import tqdm
  from dgl import DGLGraph, nn
  from dgl.distributed import DistDataLoader
  from ogb.nodeproppred import DglNodePropPredDataset
  from torch.multiprocessing import Queue
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data import DataLoader
label_desc: Use image node_id 4 for calling a global function with example usage: run(args, device, data) without return types