docling-project/docling · error · RuntimeError

Expected ONNX logits output shape [batch_size, num_classes],

Error message

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

What it means

The ONNX model's first output tensor is not 2-D. The engine requires [batch_size, num_classes] logits to softmax per row; any other rank (flat vector, extra dimension, scalar) fails here. This means the loaded ONNX export's output layout does not match the image-classification contract.

Source

Thrown at docling/models/inference_engines/image_classification/onnxruntime_engine.py:186

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

        output_tensors = self._session.run(
            [self._output_name],
            {
                self._input_name: input_tensor,
            },
        )

        if len(output_tensors) < 1:
            raise RuntimeError(
                "Expected ONNX model to return at least 1 output containing logits"
            )

        logits_batch = np.asarray(output_tensors[0], dtype=np.float32)
        if logits_batch.ndim != 2:
            raise RuntimeError(
                "Expected ONNX 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,
        )

View on GitHub (pinned to 61d76f1ff3)

Solutions

1. Log logits_batch.shape and inspect the ONNX graph's output shape (session.get_outputs()[0].shape) to identify the extra/missing dimension.
2. Re-export the model so logits are [batch, num_classes] (keep the batch axis explicit; remove wrapper dims).
3. Point options.model_filename at the classification-head ONNX file if the repo ships multiple exports.
4. Ensure the preprocessing (processor config) matches the model so pixel_values produce the expected batch axis.

Example fix

# before: export squeezes batch dim -> output [C]
torch.onnx.export(model, x, ...)

# after: keep batch dim so output is [N, C]
torch.onnx.export(model, x, ..., dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}})

import onnxruntime as ort

sess = ort.InferenceSession(str(model_path))
out_shape = sess.get_outputs()[0].shape
if out_shape and len(out_shape) != 2:
    raise ValueError(f"model output shape {out_shape} is not [batch, classes]; wrong export?")

try:
    engine.predict_batch(batch)
except RuntimeError as e:
    if "logits output shape" in str(e):
        log.error("ONNX export has wrong output rank: %s", e)
        raise  # requires model re-export; not retryable
    raise

When it happens

Trigger: OnnxRuntimeImageClassificationEngine.predict_batch() when np.asarray(output_tensors[0]).ndim != 2 — e.g. a model exported with output [1, N, C], a squeezed [C] for batch=1, or a non-classification model loaded by mistake.

Common situations: Exporting a model with the classifier head wrapped in extra ops; ONNX exports with fixed batch dim collapsing single-item batches; pointing model_filename at an object-detection or feature model; older exports with different head conventions.

Related errors

• Expected logits output shape [batch_size, num_classes], got
• Expected OnnxRuntimeImageClassificationEngineOptions, got {t
• ONNX model file '{model_filename}' not found: {model_path}
• ONNX model exposes no inputs
• ONNX model exposes no outputs

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