huggingface/pytorch-image-models · warning

It is highly recommended to have 'opt_einsum' installed for

Error message

It is highly recommended to have 'opt_einsum' installed for this optimizer.

What it means

The Kron optimizer (Kron / AdamW-style Kronecker-factored optimizer) uses torch.einsum-heavy update rules that are dramatically faster and lower-memory with the opt_einsum package's path optimization. If opt_einsum isn't installed, timm warns that you'll get slower/larger-footprint training.

Source

Thrown at timm/optim/kron.py:134

        preconditioner_update_probability: Optional[Union[Callable, float]] = None,
        max_size_triangular: int = 2048,
        min_ndim_triangular: int = 2,
        memory_save_mode: Optional[str] = None,
        momentum_into_precond_update: bool = True,
        precond_lr: float = 0.1,
        precond_init_scale: float = 1.0,
        mu_dtype: Optional[torch.dtype] = None,
        precond_dtype: Optional[torch.dtype] = None,
        decoupled_decay: bool = False,
        corrected_weight_decay: bool = False,
        flatten: bool = False,
        flatten_start_dim: int = 2,
        flatten_end_dim: int = -1,
        stochastic_weight_decay: bool = False,
        deterministic: bool = False,
    ):
        if not has_opt_einsum:
            warnings.warn("It is highly recommended to have 'opt_einsum' installed for this optimizer.")

        _validate_scalar("learning rate", lr)
        if not 0.0 <= momentum < 1.0:
            raise ValueError(f"Invalid beta parameter: {momentum}")
        if not 0.0 <= weight_decay:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")

        defaults = dict(
            lr=lr,
            momentum=momentum,
            weight_decay=weight_decay,
            preconditioner_update_probability=preconditioner_update_probability,
            max_size_triangular=max_size_triangular,
            min_ndim_triangular=min_ndim_triangular,
            memory_save_mode=memory_save_mode,
            momentum_into_precond_update=momentum_into_precond_update,
            precond_lr=precond_lr,
            precond_init_scale=precond_init_scale,

View on GitHub (pinned to 9a5261e31b)

Solutions

  1. pip install opt_einsum in the training environment
  2. If OOMing/slow before installing, also try model.set_buffer(keep_weight_decay=...) options or reduce batch size — but the install is the fix
  3. Verify with python -c "import opt_einsum" after installing

Example fix

# before (shell)
python train.py --opt kron   # warns, slow
# after
pip install opt_einsum
python train.py --opt kron
Defensive patterns

Strategy: validation

Validate before calling

try:
    import opt_einsum  # noqa: F401
    has = True
except ImportError:
    has = False
if not has:
    raise RuntimeError('pip install opt_einsum before using the Kron optimizer')

Prevention

When it happens

Trigger: timm.optim.Kron(...) or create_optimizer_v2(..., opt='kron') when the opt_einsum package is not importable (has_opt_einsum False).

Common situations: Training environments that never installed opt_einsum; slim Docker images; users switching from AdamW to Kron hitting slow steps or higher memory and OOMs.

Related errors


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