//  [python]
def calculate_loss_metrics(self, labels, outputs):
    """Calculates loss and metrics.

    Args:
      labels: Ground truth from dataset.
      outputs: fused outputs from teacher model and student model.
    Returns:
      total loss value.
    """
    if self.mode == DistillationMode.LAYER_WISE:
      teacher_feature = outputs['teacher_output_feature']
      student_feature = outputs['student_output_feature']
      feature_transfer_loss = tf_keras.losses.mean_squared_error(
          self._layer_norm(teacher_feature), self._layer_norm(student_feature))
      # feature_transfer_loss = tf.reduce_mean(feature_transfer_loss)
      feature_transfer_loss *= self.layer_wise_distill_config.hidden_distill_factor
      beta_loss, gamma_loss = _get_distribution_losses(teacher_feature,
                                                       student_feature)
      beta_loss *= self.layer_wise_distill_config.beta_distill_factor
      gamma_loss *= self.layer_wise_distill_config.gamma_distill_factor
      total_loss = feature_transfer_loss + beta_loss + gamma_loss

      teacher_attention = outputs['teacher_attention_score']
      student_attention = outputs['student_attention_score']
      attention_loss = _get_attention_loss(teacher_attention, student_attention)
      attention_loss *= self.layer_wise_distill_config.attention_distill_factor
      total_loss += attention_loss
      total_loss /= tf.cast((self.stage + 1), tf.float32)
    elif self.mode == DistillationMode.END2END:
      lm_label = labels['masked_lm_ids']
      # Shape: [batch, max_predictions_per_seq, word_vocab_size]
      lm_label = tf.one_hot(indices=lm_label,
                            depth=self.word_vocab_size,
                            on_value=1.0,
                            off_value=0.0,
                            axis=-1,
                            dtype=tf.float32)
      lm_label_weights = labels['masked_lm_weights']
      teacher_mlm_logits = outputs['teacher_pretrainer_outputs']['mlm_logits']
      teacher_labels = tf.nn.softmax(teacher_mlm_logits, axis=-1)
      gt_label = self.distill_gt_ratio * lm_label
      teacher_label = (1 - self.distill_gt_ratio) * teacher_labels
      lm_label = gt_label + teacher_label

      student_pretrainer_output = outputs['student_pretrainer_outputs']
      # Shape: [batch, max_predictions_per_seq, word_vocab_size]
      student_lm_log_probs = tf.nn.log_softmax(
          student_pretrainer_output['mlm_logits'], axis=-1)

      # Shape: [batch * max_predictions_per_seq]
      per_example_loss = tf.reshape(
          -tf.reduce_sum(student_lm_log_probs * lm_label, axis=[-1]), [-1])

      lm_label_weights = tf.reshape(labels['masked_lm_weights'], [-1])
      lm_numerator_loss = tf.reduce_sum(per_example_loss * lm_label_weights)
      lm_denominator_loss = tf.reduce_sum(lm_label_weights)
      mlm_loss = tf.math.divide_no_nan(lm_numerator_loss, lm_denominator_loss)
      total_loss = mlm_loss

      sentence_labels = labels['next_sentence_labels']
      sentence_outputs = tf.cast(
          student_pretrainer_output['next_sentence'], dtype=tf.float32)
      sentence_loss = tf.reduce_mean(
          tf_keras.losses.sparse_categorical_crossentropy(
              sentence_labels, sentence_outputs, from_logits=True))
      total_loss += sentence_loss
    else:
      raise ValueError('Training mode has to be LAYER-WISE or END2END.')

    if self.mode == DistillationMode.LAYER_WISE:
      self.train_metrics['feature_transfer_mse'].update_state(
          feature_transfer_loss)
      self.train_metrics['beta_transfer_loss'].update_state(beta_loss)
      self.train_metrics['gamma_transfer_loss'].update_state(gamma_loss)
      self.train_metrics['attention_transfer_loss'].update_state(attention_loss)
    elif self.mode == DistillationMode.END2END:
      self.train_metrics['lm_example_loss'].update_state(mlm_loss)
      self.train_metrics['next_sentence_loss'].update_state(sentence_loss)
    self.train_metrics['total_loss'].update_state(total_loss)

    return total_loss