class_docstr: BERT model with a QA (token classification) head 
class_docstr_tok: ["BERT","model","with","a","QA","(","token","classification",")","head"]
code: def forward(self, input_ids, attention_mask, token_type_ids):
      with torch.cuda.amp.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 collections
  from typing import List, Optional
  import numpy
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from transformers.models.bert.tokenization_bert import BasicTokenizer
  from nemo.collections.common.losses import SpanningLoss
  from nemo.collections.common.parts.utils import _compute_softmax
  from nemo.collections.nlp.data.question_answering.data_processor.qa_processing import QAProcessor
  from nemo.collections.nlp.data.question_answering.dataset.qa_bert_dataset import BERTQADataset
  from nemo.collections.nlp.metrics.qa_metrics import QAMetrics
  from nemo.collections.nlp.models.question_answering.qa_base_model import BaseQAModel
  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 7 for calling the BERTQAModel obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, token_type_ids) and returns: logits