code: def run(proc_id, n_gpus, args, devices, data):
      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
          torch.distributed.init_process_group(
              backend="gloo",
              init_method=dist_init_method,
              world_size=world_size,
              rank=proc_id,
              timeout=datetime.timedelta(seconds=100),
          )
      torch.cuda.set_device(dev_id)
  
      (
          g,
          train_pairs,
          index2word,
          edge_types,
          num_nodes,
          edge_type_count,
      ) = data
  
      epochs = args.epoch
      batch_size = args.batch_size
      embedding_size = args.dimensions
      embedding_u_size = args.edge_dim
      u_num = edge_type_count
      num_sampled = args.negative_samples
      dim_a = args.att_dim
      att_head = 1
      neighbor_samples = args.neighbor_samples
      num_workers = args.workers
  
      neighbor_sampler = NeighborSampler(g, [neighbor_samples])
      if n_gpus > 1:
          train_sampler = (
              torch.utils.data.distributed.DistributedSampler(
                  train_pairs,
                  num_replicas=world_size,
                  rank=proc_id,
                  shuffle=True,
                  drop_last=False,
              )
          )
          train_dataloader = torch.utils.data.DataLoader(
              train_pairs,
              batch_size=batch_size,
              collate_fn=neighbor_sampler.sample,
              num_workers=num_workers,
              sampler=train_sampler,
              pin_memory=True,
          )
      else:
          train_dataloader = torch.utils.data.DataLoader(
              train_pairs,
              batch_size=batch_size,
              collate_fn=neighbor_sampler.sample,
              num_workers=num_workers,
              shuffle=True,
              drop_last=False,
              pin_memory=True,
          )
  
      model = DGLGATNE(
          num_nodes,
          embedding_size,
          embedding_u_size,
          edge_types,
          edge_type_count,
          dim_a,
      )
  
      nsloss = NSLoss(num_nodes, num_sampled, embedding_size)
  
      model.to(dev_id)
      if n_gpus > 1:
          model = DistributedDataParallel(
              model, device_ids=[dev_id], output_device=dev_id
          )
  
      nsloss.to(dev_id)
  
      if n_gpus > 1:
          mmodel = model.module
      else:
          mmodel = model
  
      embeddings_params = list(
          map(id, mmodel.node_embeddings.parameters())
      ) + list(map(id, mmodel.node_type_embeddings.parameters()))
      weights_params = list(map(id, nsloss.weights.parameters()))
  
      optimizer = torch.optim.Adam(
          [
              {
                  "params": filter(
                      lambda p: id(p) not in embeddings_params,
                      model.parameters(),
                  )
              },
              {
                  "params": filter(
                      lambda p: id(p) not in weights_params,
                      nsloss.parameters(),
                  )
              },
          ],
          lr=2e-3,
      )
  
      sparse_optimizer = torch.optim.SparseAdam(
          [
              {"params": mmodel.node_embeddings.parameters()},
              {"params": mmodel.node_type_embeddings.parameters()},
              {"params": nsloss.weights.parameters()},
          ],
          lr=2e-3,
      )
  
      if n_gpus > 1:
          torch.distributed.barrier()
  
      if proc_id == 0:
          start = time.time()
  
      for epoch in range(epochs):
          if n_gpus > 1:
              train_sampler.set_epoch(epoch)
          model.train()
  
          data_iter = train_dataloader
          if proc_id == 0:
              data_iter = tqdm(
                  train_dataloader,
                  desc="epoch %d" % (epoch),
                  total=(len(train_pairs) + (batch_size - 1))
                  // batch_size,
              )
              avg_loss = 0.0
  
          for i, (block, head_invmap, tails, block_types) in enumerate(
              data_iter
          ):
              optimizer.zero_grad()
              sparse_optimizer.zero_grad()
              # embs: [batch_size, edge_type_count, embedding_size]
              block_types = block_types.to(dev_id)
              embs = model(block[0].to(dev_id))[head_invmap]
              embs = embs.gather(
                  1,
                  block_types.view(-1, 1, 1).expand(
                      embs.shape[0], 1, embs.shape[2]
                  ),
              )[:, 0]
              loss = nsloss(
                  block[0].dstdata[dgl.NID][head_invmap].to(dev_id),
                  embs,
                  tails.to(dev_id),
              )
              loss.backward()
              optimizer.step()
              sparse_optimizer.step()
  
              if proc_id == 0:
                  avg_loss += loss.item()
  
                  post_fix = {
                      "avg_loss": avg_loss / (i + 1),
                      "loss": loss.item(),
                  }
                  data_iter.set_postfix(post_fix)
  
          if n_gpus > 1:
              torch.distributed.barrier()
  
          if proc_id == 0:
              model.eval()
              # {'1': {}, '2': {}}
              final_model = dict(
                  zip(
                      edge_types,
                      [dict() for _ in range(edge_type_count)],
                  )
              )
              for i in range(num_nodes):
                  train_inputs = (
                      torch.tensor([i for _ in range(edge_type_count)])
                      .unsqueeze(1)
                      .to(dev_id)
                  )  # [i, i]
                  train_types = (
                      torch.tensor(list(range(edge_type_count)))
                      .unsqueeze(1)
                      .to(dev_id)
                  )  # [0, 1]
                  pairs = torch.cat(
                      (train_inputs, train_inputs, train_types), dim=1
                  )  # (2, 3)
                  (
                      train_blocks,
                      train_invmap,
                      fake_tails,
                      train_types,
                  ) = neighbor_sampler.sample(pairs.cpu())
  
                  node_emb = model(train_blocks[0].to(dev_id))[
                      train_invmap
                  ]
                  node_emb = node_emb.gather(
                      1,
                      train_types.to(dev_id)
                      .view(-1, 1, 1)
                      .expand(node_emb.shape[0], 1, node_emb.shape[2]),
                  )[:, 0]
  
                  for j in range(edge_type_count):
                      final_model[edge_types[j]][index2word[i]] = (
                          node_emb[j].cpu().detach().numpy()
                      )
  
              valid_aucs, valid_f1s, valid_prs = [], [], []
              test_aucs, test_f1s, test_prs = [], [], []
              for i in range(edge_type_count):
                  if args.eval_type == "all" or edge_types[
                      i
                  ] in args.eval_type.split(","):
                      tmp_auc, tmp_f1, tmp_pr = evaluate(
                          final_model[edge_types[i]],
                          valid_true_data_by_edge[edge_types[i]],
                          valid_false_data_by_edge[edge_types[i]],
                          num_workers,
                      )
                      valid_aucs.append(tmp_auc)
                      valid_f1s.append(tmp_f1)
                      valid_prs.append(tmp_pr)
  
                      tmp_auc, tmp_f1, tmp_pr = evaluate(
                          final_model[edge_types[i]],
                          testing_true_data_by_edge[edge_types[i]],
                          testing_false_data_by_edge[edge_types[i]],
                          num_workers,
                      )
                      test_aucs.append(tmp_auc)
                      test_f1s.append(tmp_f1)
                      test_prs.append(tmp_pr)
              print("valid auc:", np.mean(valid_aucs))
              print("valid pr:", np.mean(valid_prs))
              print("valid f1:", np.mean(valid_f1s))
  
      if proc_id == 0:
          end = time.time()
          average_auc = np.mean(test_aucs)
          average_f1 = np.mean(test_f1s)
          average_pr = np.mean(test_prs)
          print("Overall ROC-AUC:", average_auc)
          print("Overall PR-AUC", average_pr)
          print("Overall F1:", average_f1)
          print("Training Time", end - start)
  
imports: import datetime
  import math
  import os
  import sys
  import time
  from collections import defaultdict
  import numpy
  import torch
  import torch.multiprocessing
  import torch.nn
  import torch.nn.functional
  from numpy import random
  from torch.nn.parallel import DistributedDataParallel
  from torch.nn.parameter import Parameter
  from tqdm.auto import tqdm
  from utils import *
  import dgl
  import dgl.function
label_desc: Use image node_id 3 for calling a global function with example usage: run(proc_id, n_gpus, args, devices, data) without return types