code: def test(model, predictor, g, x, split_edge, evaluator, batch_size):
      model.eval()
      predictor.eval()
  
      h = model(g, x)
  
      pos_train_edge = split_edge["eval_train"]["edge"].to(h.device)
      pos_valid_edge = split_edge["valid"]["edge"].to(h.device)
      neg_valid_edge = split_edge["valid"]["edge_neg"].to(h.device)
      pos_test_edge = split_edge["test"]["edge"].to(h.device)
      neg_test_edge = split_edge["test"]["edge_neg"].to(h.device)
  
      def get_pred(test_edges, h):
          preds = []
          for perm in DataLoader(range(test_edges.size(0)), batch_size):
              edge = test_edges[perm].t()
              preds += [
                  predictor(h[edge[0]], h[edge[1]]).squeeze().cpu()
              ]
          pred = torch.cat(preds, dim=0)
          return pred
  
      pos_train_pred = get_pred(pos_train_edge, h)
      pos_valid_pred = get_pred(pos_valid_edge, h)
      neg_valid_pred = get_pred(neg_valid_edge, h)
      pos_test_pred = get_pred(pos_test_edge, h)
      neg_test_pred = get_pred(neg_test_edge, h)
  
      results = {}
      for K in [20, 50, 100]:
          evaluator.K = K
          train_hits = evaluator.eval(
              {
                  "y_pred_pos": pos_train_pred,
                  "y_pred_neg": neg_valid_pred,
              }
          )[f"hits@{K}"]
          valid_hits = evaluator.eval(
              {
                  "y_pred_pos": pos_valid_pred,
                  "y_pred_neg": neg_valid_pred,
              }
          )[f"hits@{K}"]
          test_hits = evaluator.eval(
              {
                  "y_pred_pos": pos_test_pred,
                  "y_pred_neg": neg_test_pred,
              }
          )[f"hits@{K}"]
  
          results[f"Hits@{K}"] = (train_hits, valid_hits, test_hits)
  
      return results
  
imports: import argparse
  import math
  import dgl
  import torch
  import torch.nn.functional
  from dgl.dataloading.negative_sampler import GlobalUniform
  from dgl.nn.pytorch import GraphConv, SAGEConv
  from ogb.linkproppred import DglLinkPropPredDataset, Evaluator
  from torch.nn import Linear
  from torch.utils.data import DataLoader
label_desc: Use image node_id 2 for calling a global function with example usage: test(model, predictor, g, x, split_edge, evaluator, batch_size) and returns: results, pred