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
- pip install opt_einsum in the training environment
- If OOMing/slow before installing, also try model.set_buffer(keep_weight_decay=...) options or reduce batch size — but the install is the fix
- 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
- Add opt_einsum to training environment requirements
- Include it in Dockerfiles/base conda envs when using timm optimizers
- Watch step time and memory after switching optimizers
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
- {name} must be a scalar or scalar tensor.
- Invalid {name}: {value}
- adamw_lr is deprecated, use fallback_lr_scale=adamw_lr/lr in
- Input image must have positive dimensions, got H={height}, W
- Invalid class map file, expected a dict ({class_map_path}).
AI-assisted analysis of huggingface/pytorch-image-models@9a5261e31b (2026-08-27).
Data as JSON: /api/errors/c00774220570a708.
Report an issue: GitHub.