docstr: Forward function that calculate the score of a batch of triples.
  In the 'single' mode, sample is a batch of triple.
  In the 'head-batch' or 'tail-batch' mode, sample consists two part.
  The first part is usually the positive sample.
  And the second part is the entities in the negative samples.
  Because negative samples and positive samples usually share two elements 
  in their triple ((head, relation) or (relation, tail)).
docstr_tok: ["Forward","function","that","calculate","the","score","of","a","batch","of","triples",".","In","the","'single","'","mode",",","sample","is","a","batch","of","triple",".","In","the","'head-batch","'","or","'tail-batch","'","mode",",","sample","consists","two","part",".","The","first","part","is","usually","the","positive","sample",".","And","the","second","part","is","the","entities","in","the","negative","samples",".","Because","negative","samples","and","positive","samples","usually","share","two","elements","in","their","triple","(","(","head",",","relation",")","or","(","relation",",","tail",")",")","."]
code: def forward(self, sample, mode="single"):
      """
      Forward function that calculate the score of a batch of triples.
      In the 'single' mode, sample is a batch of triple.
      In the 'head-batch' or 'tail-batch' mode, sample consists two part.
      The first part is usually the positive sample.
      And the second part is the entities in the negative samples.
      Because negative samples and positive samples usually share two elements
      in their triple ((head, relation) or (relation, tail)).
      """
  
      if mode == "single":
          batch_size, negative_sample_size = sample.size(0), 1
  
          head = torch.index_select(
              self.entity_embedding, dim=0, index=sample[:, 0]
          ).unsqueeze(1)
  
          relation = torch.index_select(
              self.relation_embedding, dim=0, index=sample[:, 1]
          ).unsqueeze(1)
  
          tail = torch.index_select(
              self.entity_embedding, dim=0, index=sample[:, 2]
          ).unsqueeze(1)
  
      elif mode == "head-batch":
          tail_part, head_part = sample
          batch_size, negative_sample_size = head_part.size(
              0
          ), head_part.size(1)
  
          head = torch.index_select(
              self.entity_embedding, dim=0, index=head_part.view(-1)
          ).view(batch_size, negative_sample_size, -1)
  
          relation = torch.index_select(
              self.relation_embedding, dim=0, index=tail_part[:, 1]
          ).unsqueeze(1)
  
          tail = torch.index_select(
              self.entity_embedding, dim=0, index=tail_part[:, 2]
          ).unsqueeze(1)
  
      elif mode == "tail-batch":
          head_part, tail_part = sample
          batch_size, negative_sample_size = tail_part.size(
              0
          ), tail_part.size(1)
  
          head = torch.index_select(
              self.entity_embedding, dim=0, index=head_part[:, 0]
          ).unsqueeze(1)
  
          relation = torch.index_select(
              self.relation_embedding, dim=0, index=head_part[:, 1]
          ).unsqueeze(1)
  
          tail = torch.index_select(
              self.entity_embedding, dim=0, index=tail_part.view(-1)
          ).view(batch_size, negative_sample_size, -1)
  
      else:
          raise ValueError("mode %s not supported" % mode)
  
      model_func = {
          "TransE": self.TransE,
          "DistMult": self.DistMult,
          "ComplEx": self.ComplEx,
          "RotatE": self.RotatE,
      }
  
      if self.model_name in model_func:
          score = model_func[self.model_name](
              head, relation, tail, mode
          )
      else:
          raise ValueError("model %s not supported" % self.model_name)
  
      return score
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import logging
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.utils.data import DataLoader
  from dataloader import TestDataset
  from collections import defaultdict
  from ogb.linkproppred import Evaluator
label_desc: Use image node_id 2 for calling the KGEModel obj's underlying member method code with example usage: obj.forward(sample, mode) and returns: score