1 files changed, 3 insertions, 4 deletions
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Shape.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Shape.java
index 361729a8c14..79f3012c327 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Shape.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Shape.java
@@ -19,11 +19,10 @@ public class Shape extends IntermediateOperation {
 
     @Override
     protected OrderedTensorType lazyGetType() {
-        if (!allInputTypesPresent(1)) {
-            return null;
-        }
+        if ( ! allInputTypesPresent(1)) return null;
+
         OrderedTensorType inputType = inputs.get(0).type().get();
-        return new OrderedTensorType.Builder()
+        return new OrderedTensorType.Builder(resultValueType())
                 .add(TensorType.Dimension.indexed(vespaName(), inputType.dimensions().size()))
                 .build();
     }