tensorflow/models · error · ValueError

Training mode has to be LAYER-WISE or END2END.

Error message

Training mode has to be LAYER-WISE or END2END.

What it means

Error "Training mode has to be LAYER-WISE or END2END." thrown in tensorflow/models.

Source

Thrown at official/projects/edgetpu/nlp/mobilebert_edgetpu_trainer.py:356

      # 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)  # pyrefly: ignore[unbound-name]
      self.train_metrics['beta_transfer_loss'].update_state(beta_loss)  # pyrefly: ignore[unbound-name]
      self.train_metrics['gamma_transfer_loss'].update_state(gamma_loss)  # pyrefly: ignore[unbound-name]
      self.train_metrics['attention_transfer_loss'].update_state(attention_loss)  # pyrefly: ignore[unbound-name]
    elif self.mode == DistillationMode.END2END:
      self.train_metrics['lm_example_loss'].update_state(mlm_loss)  # pyrefly: ignore[unbound-name]
      self.train_metrics['next_sentence_loss'].update_state(sentence_loss)  # pyrefly: ignore[unbound-name]
    self.train_metrics['total_loss'].update_state(total_loss)

    return total_loss

  def calculate_accuracy_metrics(self, labels, outputs, metrics):
    """Calculates metrics that are not related to the losses."""
    if self.mode == DistillationMode.END2END:
      student_pretrainer_output = outputs['student_pretrainer_outputs']

View on GitHub (pinned to e006f5f0d5)

When it happens

Trigger: Thrown at official/projects/edgetpu/nlp/mobilebert_edgetpu_trainer.py:356 when the library encounters an invalid state.

Common situations: See trigger scenarios.

AI-assisted analysis of tensorflow/models@e006f5f0d5 (2026-08-24). Data as JSON: /api/errors/ddfce874cce400f7. Report an issue: GitHub.