code: def train(
      world_size,
      rank,
      graph,
      features,
      train_set,
      valid_set,
      num_classes,
      model,
      device,
  ):
      optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
      # Create training data loader.
      dataloader = create_dataloader(
          graph,
          features,
          train_set,
          device,
          drop_last=False,
          shuffle=True,
          drop_uneven_inputs=False,
      )
  
      for epoch in range(5):
          epoch_start = time.time()
  
          model.train()
          total_loss = torch.tensor(0, dtype=torch.float).to(device)
          with Join([model]):
              for step, data in (
                  tqdm.tqdm(enumerate(dataloader))
                  if rank == 0
                  else enumerate(dataloader)
              ):
                  # The input features are from the source nodes in the first
                  # layer's computation graph.
                  x = data.node_features["feat"]
  
                  # The ground truth labels are from the destination nodes
                  # in the last layer's computation graph.
                  y = data.labels
  
                  blocks = data.blocks
  
                  y_hat = model(blocks, x)
  
                  # Compute loss.
                  loss = F.cross_entropy(y_hat, y)
  
                  optimizer.zero_grad()
                  loss.backward()
                  optimizer.step()
  
                  total_loss += loss
  
          # Evaluate the model.
          if rank == 0:
              print("Validating...")
          acc = (
              evaluate(
                  rank,
                  model,
                  graph,
                  features,
                  valid_set,
                  num_classes,
                  device,
              )
              / world_size
          )
          ########################################################################
          # (HIGHLIGHT) Collect accuracy and loss values from sub-processes and
          # obtain overall average values.
          #
          # `torch.distributed.reduce` is used to reduce tensors from all the
          # sub-processes to a specified process, ReduceOp.SUM is used by default.
          ########################################################################
          dist.reduce(tensor=acc, dst=0)
          total_loss /= step + 1
          dist.reduce(tensor=total_loss, dst=0)
          dist.barrier()
  
          epoch_end = time.time()
          if rank == 0:
              print(
                  f"Epoch {epoch:05d} | "
                  f"Average Loss {total_loss.item() / world_size:.4f} | "
                  f"Accuracy {acc.item():.4f} | "
                  f"Time {epoch_end - epoch_start:.4f}"
              )
  
imports: import os
  import time
  import dgl.graphbolt
  import dgl.nn
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  import torch.nn.functional
  import torchmetrics.functional
  import tqdm
  from torch.distributed.algorithms.join import Join
  from torch.nn.parallel import DistributedDataParallel
label_desc: Use image node_id 3 for calling a global function with example usage: train(world_size, rank, graph, features, train_set, valid_set, num_classes, model, device) without return types