tensorflow/models · error · ValueError

Swapping weights must occur under a tf.distribute.Strategy.

Error message

Swapping weights must occur under a tf.distribute.Strategy.

What it means

Error "Swapping weights must occur under a tf.distribute.Strategy." thrown in tensorflow/models.

Source

Thrown at official/modeling/optimization/ema_optimizer.py:176

    # Update moving average with the latest value.
    for average, normal in zip(self._average_weights, self._model_weights):  # pyrefly: ignore[bad-argument-type]
      strategy.extended.update(
          average, _apply_moving, args=(normal,), group=False
      )

  def swap_weights(self):
    """Swap the average and moving weights.

    This is a convenience method to allow one to evaluate the averaged weights
    at test time. Loads the weights stored in `self._average` into the model,
    keeping a copy of the original model weights. Swapping twice will return
    the original weights.
    """
    if tf.distribute.in_cross_replica_context():
      strategy = tf.distribute.get_strategy()
      strategy.run(self._swap_weights, args=())
    else:
      raise ValueError(
          'Swapping weights must occur under a tf.distribute.Strategy.'
      )

  @tf.function
  def _swap_weights(self):
    def fn_0(a, b):
      a.assign_add(b)
      return a
    def fn_1(b, a):
      b.assign(a - b)
      return b
    def fn_2(a, b):
      a.assign_sub(b)
      return a

    def _swap(strategy, a_and_b):
      """Swap `a` and `b` and mirror to all devices."""
      for a, b in a_and_b:

View on GitHub (pinned to e006f5f0d5)

When it happens

Trigger: Thrown at official/modeling/optimization/ema_optimizer.py:176 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/d0da30ff5d1a127e. Report an issue: GitHub.