tensorflow/models · error · ValueError

steps should be an Tensor. Python object may cause…

Error message

steps should be an Tensor. Python object may cause retracing.

What it means

Type guard inside the tf.function-decorated train_steps helper: the steps argument must be a tf.Tensor, not a Python int. Python values become symbolic constants baked into the tf.function trace, so varying step counts would trigger a costly retracing of the whole distributed training graph on every call.

Solutions

1. Convert the step count before calling: train_steps(iterator, tf.convert_to_tensor(num_steps, dtype=tf.int32))
2. Pass a tf.constant/ tf.range-derived tensor rather than a plain Python integer

When it happens

Trigger: Thrown at official/legacy/bert/model_training_utils.py:369 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/f0af70b984d5d800. Report an issue: GitHub.

Appendix: source

Thrown at official/legacy/bert/model_training_utils.py:369

      # For reporting, the metric takes the mean of losses.
      train_loss_metric.update_state(raw_loss)
      for metric in train_metrics:
        metric.update_state(labels, model_outputs)

    @tf.function
    def train_steps(iterator, steps):
      """Performs distributed training steps in a loop.

      Args:
        iterator: the distributed iterator of training datasets.
        steps: an tf.int32 integer tensor to specify number of steps to run
          inside host training loop.

      Raises:
        ValueError: Any of the arguments or tensor shapes are invalid.
      """
      if not isinstance(steps, tf.Tensor):
        raise ValueError('steps should be an Tensor. Python object may cause '
                         'retracing.')

      for _ in tf.range(steps):
        strategy.run(_replicated_step, args=(next(iterator),))  # pyrefly: ignore[missing-attribute]

    def train_single_step(iterator):
      """Performs a distributed training step.

      Args:
        iterator: the distributed iterator of training datasets.

      Raises:
        ValueError: Any of the arguments or tensor shapes are invalid.
      """
      strategy.run(_replicated_step, args=(next(iterator),))  # pyrefly: ignore[missing-attribute]

    def test_step(iterator):
      """Calculates evaluation metrics on distributed devices."""

View on GitHub (pinned to e006f5f0d5)