def postprocess_fn(batch: Dict[str, Any],
                       auxiliary_output: Dict[str, Any]) -> Dict[str, Any]:
      """Function that prepares model's input and output for serialization."""

      new_auxiliary_output = {}
      new_auxiliary_output.update(auxiliary_output)
      encoder_specific_features = encoder_postprocess_fn(
          batch, new_auxiliary_output)
      new_auxiliary_output.update(encoder_specific_features)
      return base_postprocess_fn(batch, new_auxiliary_output)
Function that prepares model's input and output for serialization.
language/mentionmemory/tasks/downstream_encoder_task.py