Expect float tensors from Onnx

1 files changed, 3 insertions, 3 deletions

diff --git a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxMnistSoftmaxImportTestCase.java b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxMnistSoftmaxImportTestCase.java
index 07814687dc6..35c853bd746 100644
--- a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxMnistSoftmaxImportTestCase.java
+++ b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxMnistSoftmaxImportTestCase.java
@@ -31,13 +31,13 @@ public class OnnxMnistSoftmaxImportTestCase {
 
         Tensor constant0 = Tensor.from(model.largeConstants().get("test_Variable"));
         assertNotNull(constant0);
-        assertEquals(new TensorType.Builder().indexed("d2", 784).indexed("d1", 10).build(),
+        assertEquals(new TensorType.Builder(TensorType.Value.FLOAT).indexed("d2", 784).indexed("d1", 10).build(),
                      constant0.type());
         assertEquals(7840, constant0.size());
 
         Tensor constant1 = Tensor.from(model.largeConstants().get("test_Variable_1"));
         assertNotNull(constant1);
-        assertEquals(new TensorType.Builder().indexed("d1", 10).build(), constant1.type());
+        assertEquals(new TensorType.Builder(TensorType.Value.FLOAT).indexed("d1", 10).build(), constant1.type());
         assertEquals(10, constant1.size());
 
         // Check inputs
@@ -52,7 +52,7 @@ public class OnnxMnistSoftmaxImportTestCase {
                      output.expression());
         assertEquals(TensorType.fromSpec("tensor<float>(d0[],d1[784])"),
                      model.inputs().get(model.defaultSignature().inputs().get("Placeholder")));
-        assertEquals("{Placeholder=tensor(d0[],d1[784])}", output.argumentTypes().toString());
+        assertEquals("{Placeholder=tensor<float>(d0[],d1[784])}", output.argumentTypes().toString());
     }
 
     @Test