From 8f1ec4820b92fd64cb2785204ce58bdb67f9132a Mon Sep 17 00:00:00 2001 From: Sebastian Larsson Date: Mon, 14 Jul 2025 10:48:05 +0200 Subject: [PATCH] Arm backend: Change node.name to buffer name when indexing state_dict Using node.name for indexing state_dict can result in the following: `KeyError: 'b__tensor_constant1'`. Using the following to fetch the buffer name is the correct way to do it: ``` buffer_name = self.exported_program.graph_signature.inputs_to_buffers[ node.name ] ``` Change-Id: Iea9a672cca921fa523788ed33adaa428035c7619 Signed-off-by: Sebastian Larsson --- backends/arm/_passes/cast_int64_pass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/arm/_passes/cast_int64_pass.py b/backends/arm/_passes/cast_int64_pass.py index 0cdd0422b61..8052c8fd2ce 100644 --- a/backends/arm/_passes/cast_int64_pass.py +++ b/backends/arm/_passes/cast_int64_pass.py @@ -47,7 +47,7 @@ def _to_int32(self, graph_module: torch.fx.GraphModule): buffer_name = self.exported_program.graph_signature.inputs_to_buffers[ node.name ] - buffer = self.exported_program.state_dict[node.name] + buffer = self.exported_program.state_dict[buffer_name] self._assert_within_int32(buffer, node) logger.warning( f"Casting buffer {node.name} from torch.int64 to torch.int32"