huggingface/pytorch-image-models · error · ValueError

Cannot infer position grid from {pos_embed_w.shape[1]} token

Error message

Cannot infer position grid from {pos_embed_w.shape[1]} tokens in {checkpoint_path}

What it means

When loading position embeddings stored as a flat token sequence (2D), the loader infers a square grid via sqrt of token count. If neither the raw count nor the count minus prefix tokens is a perfect square, the grid cannot be inferred.

Source

Thrown at timm/models/vision_transformer.py:1559

            f'{prefix}pos_embedding' if big_vision
            else f'{prefix}Transformer/posembed_input/pos_embedding')
        if embeds.pos_embed is not None and pos_embed_key in w:
            pos_embed_w = _n2p(w[pos_embed_key], t=False)
            prefix_pos_embed = None
            if pos_embed_w.ndim == 2:
                pos_embed_w = pos_embed_w.unsqueeze(0)
            if pos_embed_w.ndim == 3:
                if pos_embed_w.shape[0] == 1:
                    # Flattened NLC tables may include class/register positions.
                    num_pos_tokens = pos_embed_w.shape[1]
                    grid_size = int(math.sqrt(num_pos_tokens))
                    if grid_size * grid_size != num_pos_tokens:
                        checkpoint_prefix_tokens = (
                            1 if f'{prefix}cls' in w else getattr(embeds, 'num_prefix_tokens', 0))
                        num_pos_tokens -= checkpoint_prefix_tokens
                        grid_size = int(math.sqrt(num_pos_tokens))
                        if grid_size * grid_size != num_pos_tokens:
                            raise ValueError(
                                f'Cannot infer position grid from {pos_embed_w.shape[1]} tokens '
                                f'in {checkpoint_path}')
                        prefix_pos_embed = pos_embed_w[:, :checkpoint_prefix_tokens]
                        pos_embed_w = pos_embed_w[:, checkpoint_prefix_tokens:]
                    pos_embed_w = pos_embed_w.reshape(1, grid_size, grid_size, pos_embed_w.shape[-1])
                else:
                    # Big Vision NaFlex stores the grid directly as HWC.
                    pos_embed_w = pos_embed_w.unsqueeze(0)
            if pos_embed_w.ndim != 4:
                raise ValueError(f'Unsupported position embedding shape in {checkpoint_path}: {pos_embed_w.shape}')

            if prefix_pos_embed is not None:
                prefix_index = 0
                if embeds.cls_token is not None and prefix_pos_embed.shape[1] > prefix_index:
                    embeds.cls_token.add_(prefix_pos_embed[:, prefix_index:prefix_index + 1])
                    prefix_index += 1
                if embeds.reg_token is not None and prefix_pos_embed.shape[1] > prefix_index:
                    num_reg_tokens = min(embeds.reg_token.shape[1], prefix_pos_embed.shape[1] - prefix_index)

View on GitHub (pinned to 9a5261e31b)

Solutions

1. Load the checkpoint into a model whose grid matches the training resolution (set img_size/patch_size so the grid is rectangular where supported)
2. Use a PyTorch-format port of the checkpoint that stores the grid directly (HWC) instead of flat tokens
3. Skip position embedding loading and retrain/interpolate manually

Example fix

# before
model = vit_base_patch16_224()
load_pretrained(model, 'bigvision_rect.npz')
# after
model = vit_base_patch16_224(img_size=(240, 256))  # match training grid
load_pretrained(model, 'bigvision_rect.npz')

import math
n = pos_embed_w.shape[1]
g = int(math.sqrt(n))
ok = g * g == n or (lambda m: (s := int(math.sqrt(m))) * s == m)(n - (1 if 'cls' in w else 0))
assert ok, 'non-square grid; use matching img_size'

try:
    load_pretrained(model, path)
except ValueError as e:
    if 'Cannot infer position grid' in str(e):
        model = rebuild_with_rect_grid(path)  # match training resolution
        load_pretrained(model, path)
    else:
        raise

When it happens

Trigger: Loading a JAX checkpoint trained on non-square grids (rectangular img_size or non-square patches) — token count like 240 (16x15) fails the square test.

Common situations: Loading Big Vision checkpoints fine-tuned at rectangular resolutions, or NaFlex-adjacent checkpoints, into the square-grid ViT loader.

Related errors

• Unsupported position embedding shape in {checkpoint_path}: {
• Unsupported patch embedding rank in {checkpoint_path}: {embe
• Patch embedding shape mismatch in {checkpoint_path}: checkpo
• Cannot initialize position embeddings without grid_size.Plea
• Input image must have positive dimensions, got H={height}, W

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