code: def predict(
      model, graph, device, data_iter, split="valid", mode="tail"
  ):
      model.eval()
      with th.no_grad():
          results = {}
          train_iter = iter(data_iter["{}_{}".format(split, mode)])
  
          for step, batch in enumerate(train_iter):
              triple, label = batch[0].to(device), batch[1].to(device)
              sub, rel, obj, label = (
                  triple[:, 0],
                  triple[:, 1],
                  triple[:, 2],
                  label,
              )
              pred = model(graph, sub, rel)
              b_range = th.arange(pred.size()[0], device=device)
              target_pred = pred[b_range, obj]
              pred = th.where(
                  label.bool(), -th.ones_like(pred) * 10000000, pred
              )
              pred[b_range, obj] = target_pred
  
              # compute metrics
              ranks = (
                  1
                  + th.argsort(
                      th.argsort(pred, dim=1, descending=True),
                      dim=1,
                      descending=False,
                  )[b_range, obj]
              )
              ranks = ranks.float()
              results["count"] = th.numel(ranks) + results.get(
                  "count", 0.0
              )
              results["mr"] = th.sum(ranks).item() + results.get(
                  "mr", 0.0
              )
              results["mrr"] = th.sum(1.0 / ranks).item() + results.get(
                  "mrr", 0.0
              )
              for k in [1, 3, 10]:
                  results["hits@{}".format(k)] = th.numel(
                      ranks[ranks <= (k)]
                  ) + results.get("hits@{}".format(k), 0.0)
  
      return results
  
file_path: dgl/examples/pytorch/compGCN/main.py
imports: import argparse
  from time import time
  import numpy
  import torch
  import torch.optim
  from data_loader import Data
  from models import CompGCN_ConvE
  from utils import in_out_norm
label_desc: Use image node_id 1 for calling a global function with example usage: predict(model, graph, device, data_iter, split, mode) and returns: results