model_info: fix generated onnx output tensor names

tum-ei-eda · May 13, 2024 · 523f92e · 523f92e
1 parent a02aa3e
commit 523f92e
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/mlonmcu/flow/tvm/backend/model_info.py b/mlonmcu/flow/tvm/backend/model_info.py
@@ -247,8 +247,11 @@ def _helper(tensors):
         out_tensors = _helper(model.graph.output)
 
         # TVM seems to ignore the original output names for ONNX models
-        for i, t in enumerate(out_tensors):
-            t.name = f"output{i}"
+        if len(out_tensors) == 1:
+            out_tensors[0].name = "output"
+        else:
+            for i, t in enumerate(out_tensors):
+                t.name = f"output{i}"
         super().__init__(in_tensors, out_tensors)