cell-type conversions should match

1 files changed, 18 insertions, 5 deletions

diff --git a/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/TensorConverter.java b/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/TensorConverter.java
index 9c79961eddf..eef75a32c0a 100644
--- a/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/TensorConverter.java
+++ b/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/TensorConverter.java
@@ -180,12 +180,25 @@ class TensorConverter {
     }
 
     static private TensorType.Value toVespaValueType(TensorInfo.OnnxTensorType onnxType) {
+        // NOTE:
+        // should match best_cell_type in onnx_wrapper.cpp
         switch (onnxType) {
-            case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: return TensorType.Value.INT8;
-            case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: return TensorType.Value.BFLOAT16;
-            case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: return TensorType.Value.FLOAT;
-            case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: return TensorType.Value.DOUBLE;
-        }
+            case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL:
+            case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8:
+                return TensorType.Value.INT8;
+
+            case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16:
+                return TensorType.Value.BFLOAT16;
+
+            case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
+            case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16:
+            case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16:
+            case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
+                return TensorType.Value.FLOAT;
+
+            case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
+                return TensorType.Value.DOUBLE;
+            }
         return TensorType.Value.DOUBLE;
     }