From 91298923a0076c1b41059efb6dad2876426e4b03 Mon Sep 17 00:00:00 2001 From: Avik Chaudhuri Date: Fri, 26 Jul 2024 12:45:15 -0700 Subject: [PATCH] immutable accessors in graph signature (#4433) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/4433 splitting ET part of D60253955 Reviewed By: guangy10, zhxchen17 Differential Revision: D60295940 fbshipit-source-id: 4ad9a661a50db9b9e9bccbc13b232416d7264a49 --- exir/passes/constant_prop_pass.py | 6 +++--- exir/program/_program.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py index 354a1b071f..6ab3abbd7b 100644 --- a/exir/passes/constant_prop_pass.py +++ b/exir/passes/constant_prop_pass.py @@ -212,11 +212,11 @@ def erase_constant_node( ) -> None: # Remove corresponding tensor from param/constants dict. signature = exported_program.graph_signature - if name := signature.inputs_to_parameters.pop(node.name, None): + if name := signature.inputs_to_parameters.get(node.name, None): exported_program.state_dict.pop(name, None) - elif name := signature.inputs_to_lifted_tensor_constants.pop(node.name, None): + elif name := signature.inputs_to_lifted_tensor_constants.get(node.name, None): exported_program.constants.pop(name, None) - elif name := signature.inputs_to_buffers.pop(node.name, None): + elif name := signature.inputs_to_buffers.get(node.name, None): exported_program.constants.pop(name, None) exported_program.state_dict.pop(name, None) diff --git a/exir/program/_program.py b/exir/program/_program.py index 6fb9eca46f..fd6253a8aa 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -225,7 +225,7 @@ def lift_constant_tensor_pass(ep): return ep graph_signature = ep.graph_signature - buffers = graph_signature.buffers + buffers = list(graph_signature.buffers) fake_mode = list(ep.graph.nodes)[0].meta["val"].fake_mode first_user_input = None