huggingface/pytorch-image-models · error · ValueError

Cannot pool {N} tokens with k={k}: N must be divisible by k^

Error message

Cannot pool {N} tokens with k={k}: N must be divisible by k^2={k_squared}. Both grid dimensions must be divisible by k.

What it means

Gemma4ViT's soft pooling (global_pool='soft') average-pools k x k blocks of patch tokens, so the token count N must be divisible by k^2 (pooling_kernel_size squared). If the sequence length does not evenly tile into k-by-k cells, pooling is geometrically undefined and ValueError is raised in _avg_pool_by_positions.

Source

Thrown at timm/models/gemma4_vit.py:579

        self.root_hidden_size = hidden_size**0.5
        self.pooling_kernel_size = pooling_kernel_size

    def _avg_pool_by_positions(
            self,
            hidden_states: torch.Tensor,
            position_ids: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """2D spatial pooling on a ``k × k`` grid (k = ``self.pooling_kernel_size``).

        ``N`` patches are binned into ``k^2``-sized cells, so the pool requires
        ``N % k^2 == 0`` (caller ensures both grid dims divide by k upstream).
        ``position_ids`` follows the Gemma4-internal ``(x, y)`` convention.
        """
        N = hidden_states.shape[1]
        k = self.pooling_kernel_size
        k_squared = k * k
        if N % k_squared != 0:
            raise ValueError(
                f"Cannot pool {N} tokens with k={k}: N must be divisible by k^2={k_squared}. "
                f"Both grid dimensions must be divisible by k."
            )
        output_length = N // k_squared

        clamped_positions = position_ids.clamp(min=0)
        max_x = clamped_positions[..., 0].max(dim=-1, keepdim=True)[0] + 1
        kernel_idxs = torch.div(clamped_positions, k, rounding_mode='floor')
        kernel_idxs = kernel_idxs[..., 0] + (max_x // k) * kernel_idxs[..., 1]

        weights = F.one_hot(kernel_idxs.long(), output_length).float() / k_squared
        output = weights.transpose(1, 2) @ hidden_states.float()
        mask = torch.logical_not((weights == 0).all(dim=1))
        return output.to(hidden_states.dtype), mask

    def forward(
            self,
            hidden_states: torch.Tensor,

View on GitHub (pinned to 9a5261e31b)

Solutions

  1. Resize/crop the input so H and W are multiples of patch_size * pooling_kernel_size (e.g. multiples of 48 for patch 16, k=3)
  2. Or switch pooling: create the model with global_pool='avg' or 'none', which does not require grid alignment
  3. Pad the image to the next valid multiple before inference

Example fix

# before
model = timm.create_model('gemma4_vit_soft', pretrained=True)
out = model(torch.randn(1, 3, 224, 224))  # 224 % 48 != 0 -> error
# after
model = timm.create_model('gemma4_vit_soft', pretrained=True)
out = model(torch.randn(1, 3, 240, 240))  # 240 divisible by 48
Defensive patterns

Strategy: validation

Validate before calling

cell = model.patch_size[0] * model.pooling_kernel_size  # e.g. 16*3=48
H, W = img.shape[-2:]
if H % cell or W % cell:
    H2, W2 = (H // cell + 1) * cell, (W // cell + 1) * cell
    img = F.pad(img, (0, W2 - W, 0, H2 - H))
out = model(img)

Try / catch

try:
    out = model(x)
except ValueError as e:
    if 'divisible by k^2' in str(e) or 'divisible' in str(e):
        x = F.interpolate(x, size=(round_up(H, 48), round_up(W, 48)))
        out = model(x)
    else:
        raise

Prevention

When it happens

Trigger: Calling forward on a Gemma4ViT with global_pool='soft' (the default for these models) where the input image's H and W are divisible by patch_size but not by patch_size * pooling_kernel_size, e.g. 224px input with patch 16 and k=3 (needs multiples of 48).

Common situations: Fine-tuning at a different resolution than pretraining (e.g. 224 instead of 384) without respecting the pooling grid; NaFlex variable-resolution inputs that happen to be off-grid; test images with odd dimensions.

Related errors

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