tensorflow/models · error · ValueError

`num_steps` should be a `tf.Tensor`. Passing a Python value

Error message

`num_steps` should be a `tf.Tensor`. Passing a Python value can cause unnecessary retracing when wrapped by `tf.function`.

What it means

Error "`num_steps` should be a `tf.Tensor`. Passing a Python value can cause unnecessary retracing when wrapped by `tf.function`." thrown in tensorflow/models.

Source

Thrown at orbit/utils/loop_fns.py:116

  Returns:
    A loop function taking required `iterator` and `num_steps` parameters. If
    called inside a `tf.function`, the loop will be converted by AutoGraph into
    a `tf.while_loop` construct. See the `loop_fn` definition below for
    additional details.
  """

  def loop_fn(iterator, num_steps):
    """Makes `num_steps` calls to `step_fn(iterator)`.

    Args:
      iterator: A nested structure of `tf.data.Iterator` or
        `DistributedIterator`.
      num_steps: The number of steps in the loop. Should be passed as a
        `tf.Tensor`. Iterating until iterator exhaustion is not supported.
    """
    if not isinstance(num_steps, tf.Tensor):
      raise ValueError(
          "`num_steps` should be a `tf.Tensor`. Passing a Python value can "
          "cause unnecessary retracing when wrapped by `tf.function`.")

    for _ in tf.range(num_steps):
      # Clear out the outer name scope so the ops created inside `tf.while_loop`
      # don't get "while/" as name prefix.
      with tf.name_scope(""):  # pyrefly: ignore[bad-instantiation]
        step_fn(iterator)

  return loop_fn


def create_tf_while_loop_fn_with_state(step_fn):
  """Creates a TF while loop function with state.

  This function is similar to `create_tf_while_loop_fn`, but allowing a `state`
  to be accumulated over multiple iterations of the loop. Note that the
  structure of the `state` cannot be changed across iterations.

View on GitHub (pinned to e006f5f0d5)

Solutions

1. Pass num_steps as a tf.Tensor (e.g. tf.constant(n)) to avoid retracing.
2. Keep passing a tf.Tensor when the loop is wrapped in tf.function.

When it happens

Trigger: Thrown at orbit/utils/loop_fns.py:116 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/883f6b55bf963203. Report an issue: GitHub.