class_docstr: BERT encoder with QA head training.
class_docstr_tok: ["BERT","encoder","with","QA","head","training","."]
code: def forward(self, input_ids, attention_mask, token_type_ids):
      with autocast():
          hidden_states = self.bert_model(
              input_ids=input_ids,
              token_type_ids=token_type_ids,
              attention_mask=attention_mask,
          )
  
          if isinstance(hidden_states, tuple):
              hidden_states = hidden_states[0]
  
          logits = self.classifier(hidden_states=hidden_states)
      return logits
  
imports: import json
  from typing import Optional
  import torch
  from omegaconf import DictConfig, OmegaConf
  from pytorch_lightning import Trainer
  from torch.cuda.amp import autocast
  from nemo.collections.common.losses import SpanningLoss
  from nemo.collections.nlp.data import SquadDataset
  from nemo.collections.nlp.data.question_answering_squad.qa_squad_processing import EVALUATION_MODE, INFERENCE_MODE, TRAINING_MODE
  from nemo.collections.nlp.models.nlp_model import NLPModel
  from nemo.collections.nlp.modules.common import TokenClassifier
  from nemo.collections.nlp.parts.utils_funcs import tensor2list
  from nemo.core.classes.common import PretrainedModelInfo, typecheck
  from nemo.utils import logging
label_desc: Use image node_id 2 for calling the QAModel obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, token_type_ids) and returns: logits