tensorflow/models · error · ClassifierError

Cannot infer pooling strategy. Head input dimension

Error message

Cannot infer pooling strategy. Head input dimension {head_input_features} does not match hidden size {hidden_size} or {2 * hidden_size}.

What it means

Sentinel in _infer_pooling_from_state_dict: the classifier infers the pooling strategy by comparing the checkpoint head's input feature dimension against the backbone hidden size — equality means CLS (single CLS token), 2x means CLS_MEAN_PATCH (concatenated CLS and mean patch tokens). The saved head.weight input dimension matches neither, so the checkpoint does not correspond to this backbone variant (different hidden size) or was saved with an unknown head architecture.

Solutions

1. Verify the checkpoint was trained with the same DINOv3 backbone variant being loaded (a ViT-S checkpoint on a ViT-L backbone mismatches hidden size)
2. Check that the correct backbone — and thus hidden_size — is configured
3. Re-train or re-export the checkpoint with a supported pooling strategy so the head dimension matches
4. Manually inspect saved_state_dict['head.weight'].shape to confirm the mismatch

When it happens

Trigger: Thrown at official/projects/waste_identification_ml/model_inference_with_tracking/sam3_dinov3_tracking_pipeline/dinov3_classifier.py:131 when the library encounters an invalid state.

Common situations: See trigger scenarios.

AI-assisted analysis of tensorflow/models@e006f5f0d5 (2026-08-24). Data as JSON: /api/errors/557f8314badbd24b. Report an issue: GitHub.

Appendix: source

Thrown at official/projects/waste_identification_ml/model_inference_with_tracking/sam3_dinov3_tracking_pipeline/dinov3_classifier.py:131

  """Infers the pooling strategy from the shape of the saved head weights.

  Args:
    saved_state_dict: The checkpoint's model state dict.
    hidden_size: Backbone hidden dimensionality.

  Returns:
    The inferred pooling strategy.

  Raises:
    ClassifierError: If the head input dimension matches neither the CLS nor
      the CLS_MEAN_PATCH expectation.
  """
  head_input_features = saved_state_dict["head.weight"].shape[1]
  if head_input_features == hidden_size:
    return PoolingStrategy.CLS
  if head_input_features == 2 * hidden_size:
    return PoolingStrategy.CLS_MEAN_PATCH
  raise ClassifierError(
      "Cannot infer pooling strategy. Head input dimension "
      f"{head_input_features} does not match hidden size {hidden_size} or "
      f"{2 * hidden_size}."
  )


def _load_checkpoint_state_dict(
    checkpoint_path: pathlib.Path | str, device: torch.device
) -> dict[str, torch.Tensor]:
  """Loads the model state dict from a checkpoint file.

  Args:
    checkpoint_path: Filesystem path to the checkpoint.
    device: Target device for `map_location`.

  Returns:
    The `model_state_dict` mapping from parameter name to tensor.

View on GitHub (pinned to e006f5f0d5)