code: def run(args, device, data):
      # Unpack data
      train_nid, val_nid, test_nid, in_feats, n_classes, g = data
      shuffle = True
      # prefetch_node_feats/prefetch_labels are not supported for DistGraph yet.
      sampler = dgl.dataloading.NeighborSampler(
          [int(fanout) for fanout in args.fan_out.split(",")]
      )
      dataloader = dgl.dataloading.DistNodeDataLoader(
          g,
          train_nid,
          sampler,
          batch_size=args.batch_size,
          shuffle=shuffle,
          drop_last=False,
      )
      # Define model and optimizer
      model = DistSAGE(
          in_feats,
          args.num_hidden,
          n_classes,
          args.num_layers,
          F.relu,
          args.dropout,
      )
      model = model.to(device)
      if not args.standalone:
          if args.num_gpus == -1:
              model = th.nn.parallel.DistributedDataParallel(model)
          else:
              model = th.nn.parallel.DistributedDataParallel(
                  model, device_ids=[device], output_device=device
              )
      loss_fcn = nn.CrossEntropyLoss()
      loss_fcn = loss_fcn.to(device)
      optimizer = optim.Adam(model.parameters(), lr=args.lr)
  
      # Training loop
      iter_tput = []
      epoch = 0
      for epoch in range(args.num_epochs):
          tic = time.time()
  
          sample_time = 0
          forward_time = 0
          backward_time = 0
          update_time = 0
          num_seeds = 0
          num_inputs = 0
          start = time.time()
          # Loop over the dataloader to sample the computation dependency graph
          # as a list of blocks.
          step_time = []
  
          with model.join():
              for step, (input_nodes, seeds, blocks) in enumerate(
                  dataloader
              ):
                  tic_step = time.time()
                  sample_time += tic_step - start
                  # fetch features/labels
                  batch_inputs, batch_labels = load_subtensor(
                      g, seeds, input_nodes, "cpu"
                  )
                  batch_labels = batch_labels.long()
                  num_seeds += len(blocks[-1].dstdata[dgl.NID])
                  num_inputs += len(blocks[0].srcdata[dgl.NID])
                  # move to target device
                  blocks = [block.to(device) for block in blocks]
                  batch_inputs = batch_inputs.to(device)
                  batch_labels = batch_labels.to(device)
                  # Compute loss and prediction
                  start = time.time()
                  batch_pred = model(blocks, batch_inputs)
                  loss = loss_fcn(batch_pred, batch_labels)
                  forward_end = time.time()
                  optimizer.zero_grad()
                  loss.backward()
                  compute_end = time.time()
                  forward_time += forward_end - start
                  backward_time += compute_end - forward_end
  
                  optimizer.step()
                  update_time += time.time() - compute_end
  
                  step_t = time.time() - tic_step
                  step_time.append(step_t)
                  iter_tput.append(
                      len(blocks[-1].dstdata[dgl.NID]) / step_t
                  )
                  if step % args.log_every == 0:
                      acc = compute_acc(batch_pred, batch_labels)
                      gpu_mem_alloc = (
                          th.cuda.max_memory_allocated() / 1000000
                          if th.cuda.is_available()
                          else 0
                      )
                      print(
                          "Part {} | Epoch {:05d} | Step {:05d} | Loss {:.4f} | "
                          "Train Acc {:.4f} | Speed (samples/sec) {:.4f} | GPU "
                          "{:.1f} MB | time {:.3f} s".format(
                              g.rank(),
                              epoch,
                              step,
                              loss.item(),
                              acc.item(),
                              np.mean(iter_tput[3:]),
                              gpu_mem_alloc,
                              np.sum(step_time[-args.log_every :]),
                          )
                      )
                  start = time.time()
  
          toc = time.time()
          print(
              "Part {}, Epoch Time(s): {:.4f}, sample+data_copy: {:.4f}, "
              "forward: {:.4f}, backward: {:.4f}, update: {:.4f}, #seeds: {}, "
              "#inputs: {}".format(
                  g.rank(),
                  toc - tic,
                  sample_time,
                  forward_time,
                  backward_time,
                  update_time,
                  num_seeds,
                  num_inputs,
              )
          )
          epoch += 1
  
          if epoch % args.eval_every == 0 and epoch != 0:
              start = time.time()
              val_acc, test_acc = evaluate(
                  model if args.standalone else model.module,
                  g,
                  g.ndata["features"],
                  g.ndata["labels"],
                  val_nid,
                  test_nid,
                  args.batch_size_eval,
                  device,
              )
              print(
                  "Part {}, Val Acc {:.4f}, Test Acc {:.4f}, time: {:.4f}".format(
                      g.rank(), val_acc, test_acc, time.time() - start
                  )
              )
  
imports: import argparse
  import socket
  import time
  from contextlib import contextmanager
  import dgl
  import dgl.nn.pytorch
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import tqdm
label_desc: Use image node_id 4 for calling a global function with example usage: run(args, device, data) without return types