src/transformers/models/albert/modeling_tf_albert.py/TFAlbertForPreTraining/serving_output
class TFAlbertForPreTraining:    def serving_output(self, output: TFAlbertForPreTrainingOutput) -> TFAlbertForPreTrainingOutput:
        hs = tf.convert_to_tensor(output.hidden_states) if self.config.output_hidden_states else None
        attns = tf.convert_to_tensor(output.attentions) if self.config.output_attentions else None

        return TFAlbertForPreTrainingOutput(
            prediction_logits=output.prediction_logits,
            sop_logits=output.sop_logits,
            hidden_states=hs,
            attentions=attns,
        )