code: def run(proc_id, n_gpus, args, devices, dataset):
      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="12345"
          )
          world_size = n_gpus
          th.distributed.init_process_group(
              backend="nccl",
              init_method=dist_init_method,
              world_size=world_size,
              rank=dev_id,
          )
      if n_gpus > 0:
          th.cuda.set_device(dev_id)
  
      train_labels = dataset.train_labels
      train_truths = dataset.train_truths
      num_edges = train_truths.shape[0]
  
      reverse_types = {
          to_etype_name(k): "rev-" + to_etype_name(k)
          for k in dataset.possible_rating_values
      }
      reverse_types.update({v: k for k, v in reverse_types.items()})
      sampler = dgl.dataloading.MultiLayerNeighborSampler(
          [None], return_eids=True
      )
      sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
      dataloader = dgl.dataloading.DataLoader(
          dataset.train_enc_graph,
          {
              to_etype_name(k): th.arange(
                  dataset.train_enc_graph.num_edges(
                      etype=to_etype_name(k)
                  )
              )
              for k in dataset.possible_rating_values
          },
          sampler,
          use_ddp=n_gpus > 1,
          batch_size=args.minibatch_size,
          shuffle=True,
          drop_last=False,
      )
  
      if proc_id == 0:
          valid_dataloader = dgl.dataloading.DataLoader(
              dataset.valid_dec_graph,
              th.arange(dataset.valid_dec_graph.num_edges()),
              sampler,
              g_sampling=dataset.valid_enc_graph,
              batch_size=args.minibatch_size,
              shuffle=False,
              drop_last=False,
          )
          test_dataloader = dgl.dataloading.DataLoader(
              dataset.test_dec_graph,
              th.arange(dataset.test_dec_graph.num_edges()),
              sampler,
              g_sampling=dataset.test_enc_graph,
              batch_size=args.minibatch_size,
              shuffle=False,
              drop_last=False,
          )
  
      nd_possible_rating_values = th.FloatTensor(
          dataset.possible_rating_values
      )
      nd_possible_rating_values = nd_possible_rating_values.to(dev_id)
  
      net = Net(args=args, dev_id=dev_id)
      net = net.to(dev_id)
      if n_gpus > 1:
          net = DistributedDataParallel(
              net, device_ids=[dev_id], output_device=dev_id
          )
      rating_loss_net = nn.CrossEntropyLoss()
      learning_rate = args.train_lr
      optimizer = get_optimizer(args.train_optimizer)(
          net.parameters(), lr=learning_rate
      )
      print("Loading network finished ...\n")
  
      ### declare the loss information
      best_valid_rmse = np.inf
      no_better_valid = 0
      best_epoch = -1
      count_rmse = 0
      count_num = 0
      count_loss = 0
      print("Start training ...")
      dur = []
      iter_idx = 1
  
      for epoch in range(1, args.train_max_epoch):
          if n_gpus > 1:
              dataloader.set_epoch(epoch)
          if epoch > 1:
              t0 = time.time()
          net.train()
          with tqdm.tqdm(dataloader) as tq:
              for step, (input_nodes, pair_graph, blocks) in enumerate(
                  tq
              ):
                  head_feat, tail_feat, blocks = load_subtensor(
                      input_nodes,
                      pair_graph,
                      blocks,
                      dataset,
                      dataset.train_enc_graph,
                  )
                  frontier = blocks[0]
                  compact_g = flatten_etypes(
                      pair_graph, dataset, "train"
                  ).to(dev_id)
                  true_relation_labels = compact_g.edata["label"]
                  true_relation_ratings = compact_g.edata["rating"]
  
                  head_feat = head_feat.to(dev_id)
                  tail_feat = tail_feat.to(dev_id)
                  frontier = frontier.to(dev_id)
  
                  pred_ratings = net(
                      compact_g,
                      frontier,
                      head_feat,
                      tail_feat,
                      dataset.possible_rating_values,
                  )
                  loss = rating_loss_net(
                      pred_ratings, true_relation_labels.to(dev_id)
                  ).mean()
                  count_loss += loss.item()
                  optimizer.zero_grad()
                  loss.backward()
                  nn.utils.clip_grad_norm_(
                      net.parameters(), args.train_grad_clip
                  )
                  optimizer.step()
  
                  if proc_id == 0 and iter_idx == 1:
                      print(
                          "Total #Param of net: %d"
                          % (torch_total_param_num(net))
                      )
  
                  real_pred_ratings = (
                      th.softmax(pred_ratings, dim=1)
                      * nd_possible_rating_values.view(1, -1)
                  ).sum(dim=1)
                  rmse = (
                      (
                          real_pred_ratings
                          - true_relation_ratings.to(dev_id)
                      )
                      ** 2
                  ).sum()
                  count_rmse += rmse.item()
                  count_num += pred_ratings.shape[0]
  
                  tq.set_postfix(
                      {
                          "loss": "{:.4f}".format(
                              count_loss / iter_idx
                          ),
                          "rmse": "{:.4f}".format(
                              count_rmse / count_num
                          ),
                      },
                      refresh=False,
                  )
  
                  iter_idx += 1
  
          if epoch > 1:
              epoch_time = time.time() - t0
              print("Epoch {} time {}".format(epoch, epoch_time))
  
          if epoch % args.train_valid_interval == 0:
              if n_gpus > 1:
                  th.distributed.barrier()
              if proc_id == 0:
                  valid_rmse = evaluate(
                      args=args,
                      dev_id=dev_id,
                      net=net,
                      dataset=dataset,
                      dataloader=valid_dataloader,
                      segment="valid",
                  )
                  logging_str = "Val RMSE={:.4f}".format(valid_rmse)
  
                  if valid_rmse < best_valid_rmse:
                      best_valid_rmse = valid_rmse
                      no_better_valid = 0
                      best_epoch = epoch
                      test_rmse = evaluate(
                          args=args,
                          dev_id=dev_id,
                          net=net,
                          dataset=dataset,
                          dataloader=test_dataloader,
                          segment="test",
                      )
                      best_test_rmse = test_rmse
                      logging_str += ", Test RMSE={:.4f}".format(
                          test_rmse
                      )
                  else:
                      no_better_valid += 1
                      if (
                          no_better_valid
                          > args.train_early_stopping_patience
                          and learning_rate <= args.train_min_lr
                      ):
                          logging.info(
                              "Early stopping threshold reached. Stop training."
                          )
                          break
                      if no_better_valid > args.train_decay_patience:
                          new_lr = max(
                              learning_rate
                              * args.train_lr_decay_factor,
                              args.train_min_lr,
                          )
                          if new_lr < learning_rate:
                              logging.info(
                                  "\tChange the LR to %g" % new_lr
                              )
                              learning_rate = new_lr
                              for p in optimizer.param_groups:
                                  p["lr"] = learning_rate
                              no_better_valid = 0
                              print("Change the LR to %g" % new_lr)
              # sync on evalution
              if n_gpus > 1:
                  th.distributed.barrier()
  
          if proc_id == 0:
              print(logging_str)
      if proc_id == 0:
          print(
              "Best epoch Idx={}, Best Valid RMSE={:.4f}, Best Test RMSE={:.4f}".format(
                  best_epoch, best_valid_rmse, best_test_rmse
              )
          )
  
imports: import argparse
  import logging
  import os, time
  import random
  import string
  import traceback
  import dgl
  import numpy
  import torch
  import torch.multiprocessing
  import torch.nn
  import tqdm
  from data import MovieLens
  from model import BiDecoder, DenseBiDecoder, GCMCLayer
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data import DataLoader
  from utils import get_activation, get_optimizer, MetricLogger, to_etype_name, torch_net_info, torch_total_param_num
label_desc: Use image node_id 5 for calling a global function with example usage: run(proc_id, n_gpus, args, devices, dataset) without return types