code: def __init__(
      self,
      model_name,
      nentity,
      nrelation,
      hidden_dim,
      gamma,
      evaluator,
      double_entity_embedding=False,
      double_relation_embedding=False,
  ):
      super(KGEModel, self).__init__()
      self.model_name = model_name
      self.nentity = nentity
      self.nrelation = nrelation
      self.hidden_dim = hidden_dim
      self.epsilon = 2.0
  
      self.gamma = nn.Parameter(
          torch.Tensor([gamma]), requires_grad=False
      )
  
      self.embedding_range = nn.Parameter(
          torch.Tensor(
              [(self.gamma.item() + self.epsilon) / hidden_dim]
          ),
          requires_grad=False,
      )
  
      self.entity_dim = (
          hidden_dim * 2 if double_entity_embedding else hidden_dim
      )
      self.relation_dim = (
          hidden_dim * 2 if double_relation_embedding else hidden_dim
      )
  
      self.entity_embedding = nn.Parameter(
          torch.zeros(nentity, self.entity_dim)
      )
      nn.init.uniform_(
          tensor=self.entity_embedding,
          a=-self.embedding_range.item(),
          b=self.embedding_range.item(),
      )
  
      self.relation_embedding = nn.Parameter(
          torch.zeros(nrelation, self.relation_dim)
      )
      nn.init.uniform_(
          tensor=self.relation_embedding,
          a=-self.embedding_range.item(),
          b=self.embedding_range.item(),
      )
  
      # Do not forget to modify this line when you add a new model in the "forward" function
      if model_name not in ["TransE", "DistMult", "ComplEx", "RotatE"]:
          raise ValueError("model %s not supported" % model_name)
  
      if model_name == "RotatE" and (
          not double_entity_embedding or double_relation_embedding
      ):
          raise ValueError(
              "RotatE should use --double_entity_embedding"
          )
  
      if model_name == "ComplEx" and (
          not double_entity_embedding or not double_relation_embedding
      ):
          raise ValueError(
              "ComplEx should use --double_entity_embedding and --double_relation_embedding"
          )
  
      self.evaluator = evaluator
  
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 1 to create a new KGEModel object from inherited base classes: nn with example: obj = KGEModel(model_name, nentity, nrelation, hidden_dim, gamma, evaluator, double_entity_embedding, double_relation_embedding)