code: def test(args, model, test_samplers, step, rank=0, mode="Test"):
      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.strict_rel_part or args.soft_rel_part:
          model.load_relation(th.device("cuda:" + str(gpu_id)))
  
      with th.no_grad():
          logs = defaultdict(list)
          answers = defaultdict(list)
          for sampler in test_samplers:
              print(sampler.num_edges, sampler.batch_size)
              for query, ans, candidate in tqdm(
                  sampler,
                  disable=not args.print_on_screen,
                  total=ceil(sampler.num_edges / sampler.batch_size),
              ):
                  log = model.forward_test_wikikg(
                      query, ans, candidate, sampler.mode, gpu_id
                  ).cpu()
                  logs[sampler.mode].append(log)
                  answers[sampler.mode].append(ans)
          print("[{}] finished {} forward".format(rank, mode))
  
          input_dict = {}
          assert len(answers) == 1
          assert "h,r->t" in answers
          if "h,r->t" in answers:
              assert "h,r->t" in logs, "h,r->t not in logs"
              input_dict["h,r->t"] = {
                  "t_correct_index": th.cat(answers["h,r->t"], 0),
                  "t_pred_top10": th.cat(logs["h,r->t"], 0),
              }
      for i in range(len(test_samplers)):
          test_samplers[i] = test_samplers[i].reset()
  
      return input_dict
  
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 4 for calling a global function with example usage: test(args, model, test_samplers, step, rank, mode) and returns: input_dict