What causes "Tensor must have at least 2 dimensions, got {tensor.ndim}" in huggingface/pytorch-image-models?

Let the optimizer route 1-D params (biases, norms) to the AdamW fallback instead of Muon

huggingface/pytorch-image-models · error · ValueError

Tensor must have at least 2 dimensions, got {tensor.ndim}

Error message

Tensor must have at least 2 dimensions, got {tensor.ndim}

What it means

reshape_for_muon reshapes conv/higher-dim weight matrices into 2D for the Newton–Schulz orthogonalization. It requires tensors with at least 2 dimensions; a 0-D scalar or 1-D vector (e.g. bias, LayerNorm gain) cannot be orthogonalized.

Source

Thrown at timm/optim/muon.py:352

        tensor: torch.Tensor,
        mode: str = "flatten",
) -> Tuple[torch.Tensor, torch.Size]:
    """Reshape high-dimensional tensor for Muon processing.

    Args:
        tensor: Input tensor of shape (out, in, *spatial)
        mode: How to handle spatial dimensions
            - "flatten": Flatten spatial into output dimension (out, in*H*W)
            - "batched": Batch over spatial positions (spatial_prod, out, in) for per-position orthogonalization

    Returns:
        Reshaped tensor and original shape for restoration
    """
    original_shape = tensor.shape
    if tensor.ndim == 2:
        return tensor, original_shape
    if tensor.ndim < 2:
        raise ValueError(f"Tensor must have at least 2 dimensions, got {tensor.ndim}")

    out_ch, in_ch = tensor.shape[:2]
    if mode == "flatten":
        # Flatten: (out, in, *spatial) -> (out, in * spatial_prod)
        return tensor.reshape(out_ch, -1), original_shape
    elif mode == "batched":
        # Batched: (out, in, *spatial) -> (spatial_prod, out, in)
        # Move spatial dimension to front so zeropower_via_newtonschulz batches over it
        reshaped = tensor.reshape(out_ch, in_ch, -1)  # (out, in, spatial_prod)
        reshaped = reshaped.permute(2, 0, 1)  # (spatial_prod, out, in)
        return reshaped, original_shape
    else:
        raise ValueError(f"Unknown mode: {mode}")


def muon(
        params: List[torch.Tensor],
        grads: List[torch.Tensor],

View on GitHub (pinned to 9a5261e31b)

Solutions

Let the optimizer route 1-D params (biases, norms) to the AdamW fallback instead of Muon
If setting per-group flags, only apply use_muon to ndim>=2 matrices
Pass a correctly shaped (out, in) matrix if calling reshape_for_muon directly

Example fix

# before
param_group = {'params': [model.fc1.weight, model.fc1.bias], 'use_muon': True}
# after
param_group = {'params': [model.fc1.weight], 'use_muon': True}
bias_group = {'params': [model.fc1.bias]}  # default fallback (AdamW)

Defensive patterns

Strategy: type-guard

Validate before calling

muon_params = [p for p in params if p.ndim >= 2]
fallback_params = [p for p in params if p.ndim < 2]

Type guard

def muon_compatible(p: torch.Tensor) -> bool:
    return p.ndim >= 2

Prevention

Never route biases/norm weights into the Muon branch
Let Muon auto-route 1-D params to fallback instead of forcing use_muon

When it happens

Trigger: A parameter tensor with ndim < 2 reaching the Muon update path — e.g. biases or norm weights routed into the Muon branch instead of the fallback AdamW branch, or manually calling muon()/reshape_for_muon on a vector.

Common situations: Misconfigured param routing (Muon optimizer normally auto-routes 1-D params to fallback); overriding use_muon/use_fallback flags incorrectly; calling the muon kernel directly on flat tensors.

Related errors

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