code: def train(proc_id, n_gpus, args, dataset, g, feats, paper_offset):
      dev_id = devices[proc_id]
      if n_gpus > 1:
          dist_init_method = "tcp://{master_ip}:{master_port}".format(
              master_ip="127.0.0.1", master_port="12346"
          )
          world_size = n_gpus
          torch.distributed.init_process_group(
              backend="nccl",
              init_method=dist_init_method,
              world_size=world_size,
              rank=proc_id,
          )
  
      torch.cuda.set_device(dev_id)
  
      print("Loading masks and labels")
      train_idx = (
          torch.LongTensor(dataset.get_idx_split("train"))
          + paper_offset
      )
      valid_idx = (
          torch.LongTensor(dataset.get_idx_split("valid"))
          + paper_offset
      )
      label = dataset.paper_label
  
      print("Initializing dataloader...")
      sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 25])
  
      train_collator = ExternalNodeCollator(
          g, train_idx, sampler, paper_offset, feats, label
      )
      valid_collator = ExternalNodeCollator(
          g, valid_idx, sampler, paper_offset, feats, label
      )
      # Necessary according to https://yangkky.github.io/2019/07/08/distributed-pytorch-tutorial.html
      train_sampler = torch.utils.data.distributed.DistributedSampler(
          train_collator.dataset,
          num_replicas=world_size,
          rank=proc_id,
          shuffle=True,
          drop_last=False,
      )
      valid_sampler = torch.utils.data.distributed.DistributedSampler(
          valid_collator.dataset,
          num_replicas=world_size,
          rank=proc_id,
          shuffle=True,
          drop_last=False,
      )
  
      train_dataloader = torch.utils.data.DataLoader(
          train_collator.dataset,
          batch_size=1024,
          collate_fn=train_collator.collate,
          num_workers=4,
          sampler=train_sampler,
      )
  
      valid_dataloader = torch.utils.data.DataLoader(
          valid_collator.dataset,
          batch_size=1024,
          collate_fn=valid_collator.collate,
          num_workers=2,
          sampler=valid_sampler,
      )
  
      print("Initializing model...")
      model = RGAT(
          dataset.num_paper_features,
          dataset.num_classes,
          1024,
          5,
          2,
          4,
          0.5,
          "paper",
      ).to(dev_id)
  
      # convert BN to SyncBatchNorm. see https://pytorch.org/docs/stable/generated/torch.nn.SyncBatchNorm.html
      model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
  
      model = DistributedDataParallel(
          model, device_ids=[dev_id], output_device=dev_id
      )
      opt = torch.optim.Adam(model.parameters(), lr=0.001)
      sched = torch.optim.lr_scheduler.StepLR(
          opt, step_size=25, gamma=0.25
      )
  
      best_acc = 0
  
      for i in range(args.epochs):
          # make shuffling work properly across multiple epochs.
          # see https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
          train_sampler.set_epoch(i)
          model.train()
          with tqdm.tqdm(train_dataloader) as tq:
              for i, (input_nodes, output_nodes, mfgs) in enumerate(tq):
                  mfgs = [g.to(dev_id) for g in mfgs]
                  x = mfgs[0].srcdata["x"]
                  y = mfgs[-1].dstdata["y"]
                  y_hat = model(mfgs, x)
                  loss = F.cross_entropy(y_hat, y)
                  opt.zero_grad()
                  loss.backward()
                  opt.step()
                  acc = (y_hat.argmax(1) == y).float().mean()
                  tq.set_postfix(
                      {
                          "loss": "%.4f" % loss.item(),
                          "acc": "%.4f" % acc.item(),
                      },
                      refresh=False,
                  )
  
          # eval in each process
          model.eval()
          correct = torch.LongTensor([0]).to(dev_id)
          total = torch.LongTensor([0]).to(dev_id)
          for i, (input_nodes, output_nodes, mfgs) in enumerate(
              tqdm.tqdm(valid_dataloader)
          ):
              with torch.no_grad():
                  mfgs = [g.to(dev_id) for g in mfgs]
                  x = mfgs[0].srcdata["x"]
                  y = mfgs[-1].dstdata["y"]
                  y_hat = model(mfgs, x)
                  correct += (y_hat.argmax(1) == y).sum().item()
                  total += y_hat.shape[0]
  
          # `reduce` data into process 0
          torch.distributed.reduce(
              correct, dst=0, op=torch.distributed.ReduceOp.SUM
          )
          torch.distributed.reduce(
              total, dst=0, op=torch.distributed.ReduceOp.SUM
          )
          acc = (correct / total).item()
  
          sched.step()
  
          # process 0 print accuracy and save model
          if proc_id == 0:
              print("Validation accuracy:", acc)
  
              if best_acc < acc:
                  best_acc = acc
                  print("Updating best model...")
                  torch.save(model.state_dict(), args.model_path)
  
imports: import argparse
  import math
  import sys
  from collections import OrderedDict
  import dgl
  import dgl.nn
  import numpy
  import torch
  import torch.multiprocessing
  import torch.nn
  import torch.nn.functional
  import tqdm
  from ogb.lsc import MAG240MDataset, MAG240MEvaluator
  from torch.nn.parallel import DistributedDataParallel
label_desc: Use image node_id 1 for calling a global function with example usage: train(proc_id, n_gpus, args, dataset, g, feats, paper_offset) without return types