class_docstr: use AutoModelForMaskedLM, and select the output by another layer in the lit model
code: def training_step(
      self, batch, batch_idx
  ):  # pylint: disable=unused-argument
      input_ids, attention_mask, token_type_ids, labels, so = batch
      input_ids = input_ids.to(self.device)
      attention_mask = attention_mask.to(self.device)
      token_type_ids = token_type_ids.to(self.device)
      labels = labels.to(self.device)
      so = so.to(self.device)
      result = self.model(
          input_ids,
          attention_mask,
          token_type_ids,
          return_dict=True,
          output_hidden_states=True,
      )
      logits = result.logits
      output_embedding = result.hidden_states[-1]
      logits = self.pvp(logits, input_ids)
      loss = self.loss_fn(
          logits, labels
      ) + self.t_lambda * self.ke_loss(output_embedding, labels, so)
      # print("Train/loss: ", loss)
      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 4 for calling the BertLitModel obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: loss