class_docstr: BERT encoder with QA head training.
class_docstr_tok: ["BERT","encoder","with","QA","head","training","."]
code: def on_validation_epoch_end(self):
      if self.trainer.testing:
          prefix = "test"
          outputs = self.test_step_outputs
      else:
          prefix = "val"
          outputs = self.validation_step_outputs
  
      avg_loss = torch.stack(
          [x[f"{prefix}_loss"] for x in outputs]
      ).mean()
  
      unique_ids = torch.cat(
          [x[f"{prefix}_tensors"]["unique_ids"] for x in outputs]
      )
      start_logits = torch.cat(
          [x[f"{prefix}_tensors"]["start_logits"] for x in outputs]
      )
      end_logits = torch.cat(
          [x[f"{prefix}_tensors"]["end_logits"] for x in outputs]
      )
  
      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)
  
      exact_match, f1, all_predictions, all_nbest = -1, -1, [], []
      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
          )
          (
              exact_match,
              f1,
              all_predictions,
              all_nbest,
          ) = eval_dataset.evaluate(
              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,
          )
  
      logging.info(f"{prefix} exact match {exact_match}")
      logging.info(f"{prefix} f1 {f1}")
  
      self.log(f"{prefix}_loss", avg_loss)
      self.log(f"{prefix}_exact_match", exact_match)
      self.log(f"{prefix}_f1", f1)
      self.validation_step_outputs.clear() if prefix == "val" else self.test_step_outputs.clear()
  
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 6 for calling the QAModel obj's underlying member method code with example usage: obj.on_validation_epoch_end() without return types