code: def train(
      args,
      model,
      train_sampler,
      valid_samplers=None,
      test_samplers=None,
      rank=0,
      rel_parts=None,
      cross_rels=None,
      barrier=None,
      client=None,
  ):
      logs = []
      for arg in vars(args):
          logging.info("{:20}:{}".format(arg, getattr(args, arg)))
  
      if len(args.gpu) > 0:
          gpu_id = (
              args.gpu[rank % len(args.gpu)]
              if args.mix_cpu_gpu and args.num_proc > 1
              else args.gpu[0]
          )
      else:
          gpu_id = -1
  
      if args.async_update:
          model.create_async_update()
      if args.strict_rel_part or args.soft_rel_part:
          model.prepare_relation(th.device("cuda:" + str(gpu_id)))
      if args.soft_rel_part:
          model.prepare_cross_rels(cross_rels)
  
      if args.encoder_model_name in ["roberta", "concat"]:
          model.transform_net = model.transform_net.to(
              th.device("cuda:" + str(gpu_id))
          )
          optimizer = th.optim.Adam(
              model.transform_net.parameters(), args.mlp_lr
          )
      else:
          optimizer = None
  
      train_start = start = time.time()
      sample_time = 0
      update_time = 0
      forward_time = 0
      backward_time = 0
      for step in range(0, args.max_step):
          # print("start step", step)
          start1 = time.time()
          pos_g, neg_g = next(train_sampler)
          sample_time += time.time() - start1
  
          if client is not None:
              model.pull_model(client, pos_g, neg_g)
  
          start1 = time.time()
          if optimizer is not None:
              optimizer.zero_grad()
          loss, log = model.forward(pos_g, neg_g, gpu_id)
          forward_time += time.time() - start1
  
          start1 = time.time()
          loss.backward()
          backward_time += time.time() - start1
  
          start1 = time.time()
          if client is not None:
              model.push_gradient(client)
          else:
              model.update(gpu_id)
          if optimizer is not None:
              optimizer.step()
  
          update_time += time.time() - start1
          logs.append(log)
  
          # force synchronize embedding across processes every X steps
          if (
              args.force_sync_interval > 0
              and (step + 1) % args.force_sync_interval == 0
          ):
              barrier.wait()
  
          if (step + 1) % args.log_interval == 0:
              if (client is not None) and (
                  client.get_machine_id() != 0
              ):
                  pass
              else:
                  for k in logs[0].keys():
                      v = sum(l[k] for l in logs) / len(logs)
                      logging.info(
                          "[proc {}][Train]({}/{}) average {}: {}".format(
                              rank, (step + 1), args.max_step, k, v
                          )
                      )
                  logs = []
                  logging.info(
                      "[proc {}][Train] {} steps take {:.3f} seconds".format(
                          rank, args.log_interval, time.time() - start
                      )
                  )
                  logging.info(
                      "[proc {}]sample: {:.3f}, forward: {:.3f}, backward: {:.3f}, update: {:.3f}".format(
                          rank,
                          sample_time,
                          forward_time,
                          backward_time,
                          update_time,
                      )
                  )
                  sample_time = 0
                  update_time = 0
                  forward_time = 0
                  backward_time = 0
                  start = time.time()
  
          # if True:
          if (
              args.valid
              and (step + 1) % args.eval_interval == 0
              and step > 1
              and valid_samplers is not None
          ):
              valid_start = time.time()
              if args.strict_rel_part or args.soft_rel_part:
                  model.writeback_relation(rank, rel_parts)
              # forced sync for validation
              if barrier is not None:
                  barrier.wait()
              logging.info(
                  "[proc {}]barrier wait in validation take {:.3f} seconds:".format(
                      rank, time.time() - valid_start
                  )
              )
              valid_start = time.time()
              if valid_samplers is not None:
                  valid_input_dict = test(
                      args,
                      model,
                      valid_samplers,
                      step,
                      rank,
                      mode="Valid",
                  )
                  th.save(
                      valid_input_dict,
                      os.path.join(
                          args.save_path,
                          "valid_{}_{}.pkl".format(rank, step),
                      ),
                  )
              if test_samplers is not None:
                  test_input_dict = test(
                      args,
                      model,
                      test_samplers,
                      step,
                      rank,
                      mode="Test",
                  )
                  th.save(
                      test_input_dict,
                      os.path.join(
                          args.save_path,
                          "test_{}_{}.pkl".format(rank, step),
                      ),
                  )
              logging.info(
                  "[proc {}]validation and test take {:.3f} seconds:".format(
                      rank, time.time() - valid_start
                  )
              )
              if args.soft_rel_part:
                  model.prepare_cross_rels(cross_rels)
              if barrier is not None:
                  barrier.wait()
  
      print(
          "proc {} takes {:.3f} seconds".format(
              rank, time.time() - train_start
          )
      )
      time.sleep(10)
      if rank == 0 and not args.no_save_emb:
          save_model(args, model, None, None)
          print("proc {} model saved".format(rank))
  
      if barrier is not None:
          barrier.wait()
      print("proc {} after barrier".format(rank))
      if args.async_update:
          model.finish_async_update()
      print("proc {} finish async update".format(rank))
      if args.strict_rel_part or args.soft_rel_part:
          model.writeback_relation(rank, rel_parts)
      print("proc {} return".format(rank))
  
imports: import torch.multiprocessing
  from torch.utils.data import DataLoader
  import torch.optim
  import torch
  from distutils.version import LooseVersion
  from .models.pytorch.tensor_models import thread_wrapped_func
  from .models import KEModel
  from .utils import save_model, get_compatible_batch_size
  import os
  import logging
  import time
  from functools import wraps
  import dgl
  from dgl.contrib import KVClient
  import dgl.backend
  from .dataloader import EvalDataset
  from .dataloader import get_dataset
  import pdb
  from collections import defaultdict
  from ogb.lsc import WikiKG90MDataset, WikiKG90MEvaluator
  from tqdm.auto import tqdm
  import pickle
  from math import ceil
label_desc: Use image node_id 3 for calling a global function with example usage: train(args, model, train_sampler, valid_samplers, test_samplers, rank, rel_parts, cross_rels, barrier, client) without return types