Verify correct output node for ONNX

1 files changed, 4 insertions, 4 deletions

diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/onnx/GraphImporter.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/onnx/GraphImporter.java
index b670eca9183..ccf3c2d8fb0 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/onnx/GraphImporter.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/onnx/GraphImporter.java
@@ -198,11 +198,11 @@ class GraphImporter {
     }
 
     private static void verifyOutputTypes(Onnx.GraphProto onnxGraph, IntermediateGraph intermediateGraph) {
-        for (String outputName : intermediateGraph.outputs(intermediateGraph.defaultSignature()).values()) {
-            IntermediateOperation operation = intermediateGraph.get(outputName);
-            Onnx.ValueInfoProto onnxNode = getOutputNode(outputName, onnxGraph);
+        for (java.util.Map.Entry<String, String> output : intermediateGraph.outputs(intermediateGraph.defaultSignature()).entrySet()) {
+            IntermediateOperation operation = intermediateGraph.get(output.getValue());
+            Onnx.ValueInfoProto onnxNode = getOutputNode(output.getKey(), onnxGraph);
             OrderedTensorType type = operation.type().orElseThrow(
-                        () -> new IllegalArgumentException("Output of '" + outputName + "' has no type."));
+                        () -> new IllegalArgumentException("Output of '" + output.getValue() + "' has no type."));
             TypeConverter.verifyType(onnxNode.getType(), type);
         }
     }