huggingface/pytorch-image-models · error · RuntimeError

`requires_grad` is not supported for `step` in differentiabl

Error message

`requires_grad` is not supported for `step` in differentiable mode

What it means

In differentiable optimization mode (differentiable=True), ADOPT computes gradients through the optimizer step; if the per-parameter step counter tensor itself has requires_grad=True the higher-order autograd graph is ill-defined, so the optimizer refuses.

Source

Thrown at timm/optim/adopt.py:178

            if len(state) == 0:
                # note(crcrpar): [special device hosting for step]
                # Deliberately host `step` on CPU if both capturable and fused are off.
                # This is because kernel launches are costly on CUDA and XLA.
                state["step"] = (
                    torch.zeros((), dtype=_get_scalar_dtype(), device=p.grad.device)
                    if group["capturable"]
                    else torch.tensor(0.0, dtype=_get_scalar_dtype())
                )
                # Exponential moving average of gradient values
                state["exp_avg"] = torch.zeros_like(p.grad, memory_format=torch.preserve_format)
                # Exponential moving average of squared gradient values
                state["exp_avg_sq"] = torch.zeros_like(p.grad, memory_format=torch.preserve_format)

            exp_avgs.append(state["exp_avg"])
            exp_avg_sqs.append(state["exp_avg_sq"])

            if group["differentiable"] and state["step"].requires_grad:
                raise RuntimeError("`requires_grad` is not supported for `step` in differentiable mode")

            # Foreach without capturable does not support a tensor lr
            if group["foreach"] and torch.is_tensor(group["lr"]) and not group["capturable"]:
                raise RuntimeError("lr as a Tensor is not supported for capturable=False and foreach=True")

            state_steps.append(state["step"])
        return has_complex

    #@_use_grad_for_differentiable  # FIXME internal context mgr, can't use
    @torch.no_grad()
    def step(self, closure=None):
        """Perform a single optimization step.

        Args:
            closure (Callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        if hasattr(self, '_accelerator_graph_capture_health_check'):

View on GitHub (pinned to 9a5261e31b)

Solutions

1. Create step counters as torch.zeros(1, dtype=torch.float, requires_grad=False), or let Adopt lazily initialize state itself
2. Detach step tensors: state['step'] = state['step'].detach() before stepping
3. Only mark exp_avg/exp_avg_sq (and params) as differentiable, not step

Example fix

# before
state['step'] = torch.zeros(1, requires_grad=True)
# after
state['step'] = torch.zeros(1, requires_grad=False)

assert not any(p in opt.state and opt.state[p]['step'].requires_grad for p in group['params'])

def step_counters_safe(opt) -> bool:
    return all(not st['step'].requires_grad for st in opt.state.values() if 'step' in st)

When it happens

Trigger: Setting differentiable=True in timm.optim.Adopt and manually creating state['step'] tensors with requires_grad=True (e.g. when reimplementing state init), then calling step().

Common situations: Hyperparameter-optimization code (meta-gradients through the optimizer) that marks every state tensor as differentiable, including step counters; porting a differentiable-optimizer implementation from torch.optim.Adam.

Related errors

• lr as a Tensor is not supported for capturable=False and for
• Tensor lr must be 1-element
• Invalid learning rate: {lr}
• Invalid epsilon value: {eps}
• Invalid beta parameter at index 0: {betas[0]}

AI-assisted analysis of huggingface/pytorch-image-models@9a5261e31b (2026-08-27). Data as JSON: /api/errors/fab7cadcc16ffc67. Report an issue: GitHub.