diff --git a/src/finn/core/onnx_exec.py b/src/finn/core/onnx_exec.py
index 6c2bf6d053d3a24a5e457cd7152964bc5fa3e8d1..0c01a48a07608dcd760447e8f569128f58d86f28 100644
--- a/src/finn/core/onnx_exec.py
+++ b/src/finn/core/onnx_exec.py
@@ -60,7 +60,11 @@ def execute_node(node, context, graph):
             inp_ctx[new_iname] = inp_ctx[old_iname]
             del inp_ctx[old_iname]
         ret = execute_onnx(model, inp_ctx, False)
-        context.update(ret)
+        # output may have been renamed in partition
+        assert len(ret) == 1
+        node_oname = node.output[0]
+        model_oname = model.graph.output[0].name
+        context[node_oname] = ret[model_oname]
     else:
         if node.domain == "finn":