{"record":{"id":"aa6d005854ecbe81","repo":"docling-project/docling","slug":"expected-onnx-logits-output-shape-batch-size-num","errorCode":null,"errorMessage":"Expected ONNX logits output shape [batch_size, num_classes], got shape={logits_batch.shape}","messagePattern":"Expected ONNX logits output shape \\[batch_size, num_classes\\], got shape=(.+?)","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"docling/models/inference_engines/image_classification/onnxruntime_engine.py","lineNumber":186,"sourceCode":"        images = [item.image.convert(\"RGB\") for item in input_batch]\n        inputs = self._processor(images=images, return_tensors=\"np\")\n        input_tensor = np.asarray(inputs[\"pixel_values\"], dtype=np.float32)\n\n        output_tensors = self._session.run(\n            [self._output_name],\n            {\n                self._input_name: input_tensor,\n            },\n        )\n\n        if len(output_tensors) < 1:\n            raise RuntimeError(\n                \"Expected ONNX model to return at least 1 output containing logits\"\n            )\n\n        logits_batch = np.asarray(output_tensors[0], dtype=np.float32)\n        if logits_batch.ndim != 2:\n            raise RuntimeError(\n                \"Expected ONNX logits output shape [batch_size, num_classes], \"\n                f\"got shape={logits_batch.shape}\"\n            )\n\n        probs_batch = self._softmax(logits_batch)\n        return self._build_batch_outputs_from_probabilities(\n            input_batch=input_batch,\n            probs_batch=probs_batch,\n        )\n","sourceCodeStart":168,"sourceCodeEnd":196,"githubUrl":"https://github.com/docling-project/docling/blob/61d76f1ff3f8428065465889f7b4577da7df704c/docling/models/inference_engines/image_classification/onnxruntime_engine.py#L168-L196","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Log logits_batch.shape and inspect the ONNX graph's output shape (session.get_outputs()[0].shape) to identify the extra/missing dimension.","Re-export the model so logits are [batch, num_classes] (keep the batch axis explicit; remove wrapper dims).","Point options.model_filename at the classification-head ONNX file if the repo ships multiple exports.","Ensure the preprocessing (processor config) matches the model so pixel_values produce the expected batch axis."],"exampleFix":"# before: export squeezes batch dim -> output [C]\ntorch.onnx.export(model, x, ...)\n\n# after: keep batch dim so output is [N, C]\ntorch.onnx.export(model, x, ..., dynamic_axes={\"input\": {0: \"batch\"}, \"logits\": {0: \"batch\"}})","handlingStrategy":"validation","validationCode":"import onnxruntime as ort\n\nsess = ort.InferenceSession(str(model_path))\nout_shape = sess.get_outputs()[0].shape\nif out_shape and len(out_shape) != 2:\n    raise ValueError(f\"model output shape {out_shape} is not [batch, classes]; wrong export?\")","typeGuard":null,"tryCatchPattern":"try:\n    engine.predict_batch(batch)\nexcept RuntimeError as e:\n    if \"logits output shape\" in str(e):\n        log.error(\"ONNX export has wrong output rank: %s\", e)\n        raise  # requires model re-export; not retryable\n    raise","preventionTips":["Check session.get_outputs()[0].shape at startup and assert rank 2.","Export ONNX with dynamic batch axes to keep the batch dimension.","Use the model spec's extra_config 'model_filename' to select the classification-head export."],"tags":["onnx","tensor-shape","logits","model-export"],"backgroundTag":null,"analyzedSha":"61d76f1ff3f8428065465889f7b4577da7df704c","analyzedAt":"2026-08-14T23:53:18.727Z","schemaVersion":2},"datasetVersion":"2026-08-15T17:31:12.345Z"}