code: def train(
      proc_id,
      device,
      g,
      target_idx,
      labels,
      num_classes,
      train_idx,
      inv_target,
      model,
  ):
      # define loss function and optimizer
      loss_fcn = nn.CrossEntropyLoss()
      optimizer = torch.optim.Adam(
          model.parameters(), lr=1e-2, weight_decay=5e-4
      )
      # construct sampler and dataloader
      sampler = MultiLayerNeighborSampler([4, 4])
      train_loader = DataLoader(
          g,
          target_idx[train_idx],
          sampler,
          device=device,
          batch_size=100,
          shuffle=True,
          use_ddp=True,
      )
      # no separate validation subset, use train index instead for validation
      val_loader = DataLoader(
          g,
          target_idx[train_idx],
          sampler,
          device=device,
          batch_size=100,
          shuffle=False,
          use_ddp=True,
      )
      for epoch in range(50):
          model.train()
          total_loss = 0
          for it, (input_nodes, output_nodes, blocks) in enumerate(
              train_loader
          ):
              output_nodes = inv_target[output_nodes]
              for block in blocks:
                  block.edata["norm"] = dgl.norm_by_dst(
                      block
                  ).unsqueeze(1)
              logits = model(blocks)
              loss = loss_fcn(logits, labels[output_nodes])
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
              total_loss += loss.item()
          # torchmetric accuracy defined as num_correct_labels / num_train_nodes
          # loc_acc_split = [loc_accuracy * loc_num_train_nodes, loc_num_train_nodes]
          loc_acc_split = evaluate(
              model, labels, num_classes, val_loader, inv_target
          ).to(device)
          dist.reduce(loc_acc_split, 0)
          if proc_id == 0:
              acc = loc_acc_split[0] / loc_acc_split[1]
              print(
                  "Epoch {:05d} | Loss {:.4f} | Val. Accuracy {:.4f} ".format(
                      epoch, total_loss / (it + 1), acc.item()
                  )
              )
  
file_path: dgl/examples/pytorch/rgcn/entity_sample_multi_gpu.py
imports: import argparse
  import os
  import dgl
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  import torch.nn.functional
  from dgl.data.rdf import AIFBDataset, AMDataset, BGSDataset, MUTAGDataset
  from dgl.dataloading import DataLoader, MultiLayerNeighborSampler
  from dgl.nn.pytorch import RelGraphConv
  from torch.nn.parallel import DistributedDataParallel
  from torchmetrics.functional import accuracy
label_desc: Use image node_id 2 for calling a global function with example usage: train(proc_id, device, g, target_idx, labels, num_classes, train_idx, inv_target, model) without return types