docling-project/docling · error · RuntimeError

Expected logits output shape [batch_size, num_classes], got

Error message

Expected logits output shape [batch_size, num_classes], got shape={logits_batch.shape}

What it means

The KServe v2 model returned a logits tensor that is not 2-dimensional. The engine expects [batch_size, num_classes] so it can apply softmax per row; any other rank (e.g. [batch, 1, num_classes] or a flat vector for a multi-image batch) fails this shape check.

Source

Thrown at docling/models/inference_engines/image_classification/api_kserve_v2_engine.py:176

        images = [item.image.convert("RGB") for item in input_batch]
        processed_inputs = self._processor(images=images, return_tensors="np")
        pixel_values = np.asarray(processed_inputs["pixel_values"])

        outputs = self._kserve_client.infer(
            inputs={self._input_name: pixel_values},
            output_names=[self._output_name],
            request_parameters=self.options.request_parameters,
        )
        try:
            logits_batch = outputs[self._output_name]
        except KeyError as exc:
            raise RuntimeError(
                f"Missing expected KServe v2 output: {self._output_name}"
            ) from exc

        logits_batch = np.asarray(logits_batch, dtype=np.float32)
        if logits_batch.ndim != 2:
            raise RuntimeError(
                "Expected logits output shape [batch_size, num_classes], "
                f"got shape={logits_batch.shape}"
            )

        probs_batch = self._softmax(logits_batch)
        return self._build_batch_outputs_from_probabilities(
            input_batch=input_batch,
            probs_batch=probs_batch,
        )

    def close(self) -> None:
        if self._kserve_client is None:
            return
        self._kserve_client.close()

    def __del__(self) -> None:
        try:
            self.close()

View on GitHub (pinned to 61d76f1ff3)

Solutions

1. Log logits_batch.shape and compare with the model's expected output; reshape server-side or pick the correct output tensor that is genuinely [N, num_classes].
2. If the server squeezes the batch axis for batch=1, ensure the served graph always emits 2-D output (e.g. keep batch dim explicit).
3. Confirm the endpoint actually serves an image-classification model with per-class logits output.
4. Match the request batch to a supported dynamic batch axis.

Example fix

# before: server output squeezed to [num_classes] for batch=1
out = logits.reshape(-1)  # on server

# after: always emit [batch, num_classes]
out = logits.reshape(1, -1) if logits.ndim == 1 else logits

probe = kserve_client.infer(inputs={input_name: dummy_batch}, output_names=[output_name])
shape = np.asarray(probe[output_name]).shape
if len(shape) != 2:
    raise ValueError(f"model output rank {len(shape)} != 2 (shape={shape}); wrong model?")

try:
    engine.predict_batch(batch)
except RuntimeError as e:
    if "Expected logits output shape" in str(e):
        log.error("server logits shape unexpected: %s", e)
        raise  # server model must be fixed; client-side retry is pointless
    raise

When it happens

Trigger: ApiKserveV2ImageClassificationEngine.predict_batch() when the decoded output tensor for the requested output name has ndim != 2 — commonly a squeezed batch dimension for single-image requests, an extra channel dim, or a model that returns non-classification output.

Common situations: Serving a model with an image-classification head wrapped in extra dims; model exported for a fixed batch size of 1 collapsing the batch axis; pointing the engine at a detection/embedding model by mistake; server post-processing that flattens or reshapes outputs.

Related errors

• Invalid BYTES data: insufficient bytes for string of length
• Expected image-classification model metadata to expose at le
• Expected image-classification model metadata to expose at le
• Missing expected KServe v2 output: {self._output_name}
• Expected ONNX logits output shape [batch_size, num_classes],

AI-assisted analysis of docling-project/docling@61d76f1ff3 (2026-08-14). Data as JSON: /api/errors/ab67712dd2cb5474. Report an issue: GitHub.