Revert D60253955: immutable accessors in graph signature

Differential Revision: D60253955 Original commit changeset: eb6eb65bf17f Original Phabricator Diff: D60253955 fbshipit-source-id: f203ef791da6f7efa40bf51a6e905eba65cb6b47
pytorch · Jul 26, 2024 · 5d3ec13 · 5d3ec13
1 parent 11407f0
commit 5d3ec13
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git 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.get(node.name, None):
+    if name := signature.inputs_to_parameters.pop(node.name, None):
         exported_program.state_dict.pop(name, None)
-    elif name := signature.inputs_to_lifted_tensor_constants.get(node.name, None):
+    elif name := signature.inputs_to_lifted_tensor_constants.pop(node.name, None):
         exported_program.constants.pop(name, None)
-    elif name := signature.inputs_to_buffers.get(node.name, None):
+    elif name := signature.inputs_to_buffers.pop(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
@@ -225,7 +225,7 @@ def lift_constant_tensor_pass(ep):
         return ep
 
     graph_signature = ep.graph_signature
-    buffers = list(graph_signature.buffers)
+    buffers = graph_signature.buffers
 
     fake_mode = list(ep.graph.nodes)[0].meta["val"].fake_mode
     first_user_input = None