class_docstr: use AutoModelForMaskedLM, and select the output by another layer in the lit model
code: def pvp(self, logits, input_ids):
      # convert the [batch_size, seq_len, vocab_size] => [batch_size, num_labels]
      #! hard coded
      _, mask_idx = (input_ids == 103).nonzero(as_tuple=True)
      bs = input_ids.shape[0]
      mask_output = logits[torch.arange(bs), mask_idx]
      assert mask_idx.shape[0] == bs, "only one mask in sequence!"
      final_output = mask_output[:, self.word2label]
  
      return final_output
  
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 10 for calling the BertLitModel obj's underlying member method code with example usage: obj.pvp(logits, input_ids) and returns: final_output