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." )