{"record":{"id":"ab67712dd2cb5474","repo":"docling-project/docling","slug":"expected-logits-output-shape-batch-size-num-clas","errorCode":null,"errorMessage":"Expected logits output shape [batch_size, num_classes], got shape={logits_batch.shape}","messagePattern":"Expected logits output shape \\[batch_size, num_classes\\], got shape=(.+?)","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"docling/models/inference_engines/image_classification/api_kserve_v2_engine.py","lineNumber":176,"sourceCode":"        images = [item.image.convert(\"RGB\") for item in input_batch]\n        processed_inputs = self._processor(images=images, return_tensors=\"np\")\n        pixel_values = np.asarray(processed_inputs[\"pixel_values\"])\n\n        outputs = self._kserve_client.infer(\n            inputs={self._input_name: pixel_values},\n            output_names=[self._output_name],\n            request_parameters=self.options.request_parameters,\n        )\n        try:\n            logits_batch = outputs[self._output_name]\n        except KeyError as exc:\n            raise RuntimeError(\n                f\"Missing expected KServe v2 output: {self._output_name}\"\n            ) from exc\n\n        logits_batch = np.asarray(logits_batch, dtype=np.float32)\n        if logits_batch.ndim != 2:\n            raise RuntimeError(\n                \"Expected 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\n    def close(self) -> None:\n        if self._kserve_client is None:\n            return\n        self._kserve_client.close()\n\n    def __del__(self) -> None:\n        try:\n            self.close()","sourceCodeStart":158,"sourceCodeEnd":194,"githubUrl":"https://github.com/docling-project/docling/blob/61d76f1ff3f8428065465889f7b4577da7df704c/docling/models/inference_engines/image_classification/api_kserve_v2_engine.py#L158-L194","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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].","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).","Confirm the endpoint actually serves an image-classification model with per-class logits output.","Match the request batch to a supported dynamic batch axis."],"exampleFix":"# before: server output squeezed to [num_classes] for batch=1\nout = logits.reshape(-1)  # on server\n\n# after: always emit [batch, num_classes]\nout = logits.reshape(1, -1) if logits.ndim == 1 else logits","handlingStrategy":"validation","validationCode":"probe = kserve_client.infer(inputs={input_name: dummy_batch}, output_names=[output_name])\nshape = np.asarray(probe[output_name]).shape\nif len(shape) != 2:\n    raise ValueError(f\"model output rank {len(shape)} != 2 (shape={shape}); wrong model?\")","typeGuard":null,"tryCatchPattern":"try:\n    engine.predict_batch(batch)\nexcept RuntimeError as e:\n    if \"Expected logits output shape\" in str(e):\n        log.error(\"server logits shape unexpected: %s\", e)\n        raise  # server model must be fixed; client-side retry is pointless\n    raise","preventionTips":["Probe the served model once at startup with a dummy batch and assert [N, C] output.","Keep a fixed contract test against the deployed endpoint in CI/CD.","Match request batch size with the model's supported batch axis."],"tags":["kserve","tensor-shape","logits","remote-inference"],"backgroundTag":null,"analyzedSha":"61d76f1ff3f8428065465889f7b4577da7df704c","analyzedAt":"2026-08-14T23:53:18.727Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}