class_docstr: BERT model with a QA (token classification) head 
class_docstr_tok: ["BERT","model","with","a","QA","(","token","classification",")","head"]
code: def on_validation_epoch_end(self):
      prefix = "test" if self.trainer.testing else "val"
  
      if prefix == "val":
          avg_loss = torch.stack(
              [
                  x[f"{prefix}_loss"]
                  for x in self.validation_step_outputs
              ]
          ).mean()
  
          unique_ids = torch.cat(
              [
                  x[f"{prefix}_tensors"]["unique_ids"]
                  for x in self.validation_step_outputs
              ]
          )
          start_logits = torch.cat(
              [
                  x[f"{prefix}_tensors"]["start_logits"]
                  for x in self.validation_step_outputs
              ]
          )
          end_logits = torch.cat(
              [
                  x[f"{prefix}_tensors"]["end_logits"]
                  for x in self.validation_step_outputs
              ]
          )
          self.validation_step_outputs.clear()  # free memory
      else:
          avg_loss = torch.stack(
              [x[f"{prefix}_loss"] for x in self.test_step_outputs]
          ).mean()
  
          unique_ids = torch.cat(
              [
                  x[f"{prefix}_tensors"]["unique_ids"]
                  for x in self.test_step_outputs
              ]
          )
          start_logits = torch.cat(
              [
                  x[f"{prefix}_tensors"]["start_logits"]
                  for x in self.test_step_outputs
              ]
          )
          end_logits = torch.cat(
              [
                  x[f"{prefix}_tensors"]["end_logits"]
                  for x in self.test_step_outputs
              ]
          )
          self.test_step_outputs.clear()  # free memory
  
      all_unique_ids = []
      all_start_logits = []
      all_end_logits = []
      if torch.distributed.is_initialized():
          world_size = torch.distributed.get_world_size()
          for ind in range(world_size):
              all_unique_ids.append(torch.empty_like(unique_ids))
              all_start_logits.append(torch.empty_like(start_logits))
              all_end_logits.append(torch.empty_like(end_logits))
          torch.distributed.all_gather(all_unique_ids, unique_ids)
          torch.distributed.all_gather(all_start_logits, start_logits)
          torch.distributed.all_gather(all_end_logits, end_logits)
      else:
          all_unique_ids.append(unique_ids)
          all_start_logits.append(start_logits)
          all_end_logits.append(end_logits)
  
      eval_results, all_predictions, all_nbest = {}, [], []
      if (
          not torch.distributed.is_initialized()
          or torch.distributed.get_rank() == 0
      ):
          unique_ids = []
          start_logits = []
          end_logits = []
          for u in all_unique_ids:
              unique_ids.extend(tensor2list(u))
          for u in all_start_logits:
              start_logits.extend(tensor2list(u))
          for u in all_end_logits:
              end_logits.extend(tensor2list(u))
  
          eval_dataset = (
              self._test_dl.dataset
              if self.trainer.testing
              else self._validation_dl.dataset
          )
          eval_results, _, _ = self.evaluate(
              eval_dataset.features,
              eval_dataset.examples,
              eval_dataset.processor,
              unique_ids=unique_ids,
              start_logits=start_logits,
              end_logits=end_logits,
              n_best_size=self._cfg.dataset.n_best_size,
              max_answer_length=self._cfg.dataset.max_answer_length,
              version_2_with_negative=self._cfg.dataset.version_2_with_negative,
              null_score_diff_threshold=self._cfg.dataset.null_score_diff_threshold,
              do_lower_case=self._cfg.dataset.do_lower_case,
          )
  
      self.log(f"{prefix}_loss", avg_loss)
      for eval_key in eval_results:
          logging.info(f"{prefix} {eval_key}: {eval_results[eval_key]}")
          self.log(f"{prefix}_{eval_key}", eval_results[eval_key])
  
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 5 for calling the BERTQAModel obj's underlying member method code with example usage: obj.on_validation_epoch_end() without return types