class_docstr: use AutoModelForMaskedLM, and select the output by another layer in the lit model
code: def ke_loss(self, logits, labels, so):
      subject_embedding = []
      object_embedding = []
      bsz = logits.shape[0]
      for i in range(bsz):
          subject_embedding.append(
              torch.mean(logits[i, so[i][0] : so[i][1]], dim=0)
          )
          object_embedding.append(
              torch.mean(logits[i, so[i][2] : so[i][3]], dim=0)
          )
  
      subject_embedding = torch.stack(subject_embedding)
      object_embedding = torch.stack(object_embedding)
      # trick , the relation ids is concated,
      relation_embedding = (
          self.cur_model.get_output_embeddings().weight[
              labels + self.label_st_id
          ]
      )
  
      loss = torch.norm(
          subject_embedding + relation_embedding - object_embedding, p=2
      )
  
      return loss
  
imports: import json
  import torch
  import torch.nn
  import torch.nn.functional
  import numpy
  from .base import BaseLitModel
  from .util import dialog_f1_eval, f1_eval, compute_f1, acc, f1_score
  from transformers.optimization import get_linear_schedule_with_warmup
  from functools import partial
label_desc: Use image node_id 11 for calling the BertLitModel obj's underlying member method code with example usage: obj.ke_loss(logits, labels, so) and returns: loss