From 23e1fcbfce8a1c99f63d97c114a6c804e2ef57b7 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 1 Aug 2024 11:51:48 -0700 Subject: [PATCH] [IR] Fix error message with remove(..., safe=True) (#1768) Previously the node printed was incorrect. --- onnxscript/ir/_core.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/onnxscript/ir/_core.py b/onnxscript/ir/_core.py index 7eeba0493..b5a29cdd4 100644 --- a/onnxscript/ir/_core.py +++ b/onnxscript/ir/_core.py @@ -1647,12 +1647,12 @@ def _check_node_safe_to_remove( raise ValueError( f"Node '{node!r}' is still an output of the graph and cannot be removed when safe=True." ) - for use, _ in output.uses(): - if use in to_remove: - continue + uses_not_to_remove = [user for user, _ in output.uses() if user not in to_remove] + if uses_not_to_remove: raise ValueError( - f"Node '{use!r}' is still being used by other nodes that are not to be " - f"removed. All of its uses: {list(output.uses())!r}" + f"Output value '{output!r}' is still being used by other nodes that are not to be " + f"removed. All of its users that is not being removed: {uses_not_to_remove!r}. " + "Please make sure these nodes are no longer using the output value." )