From 04f98c085667624f7fdd6e366254dffe57acf530 Mon Sep 17 00:00:00 2001 From: Oscar Andersson Date: Thu, 8 May 2025 12:47:54 +0200 Subject: [PATCH] Update partitioner de-tagging iteration order When de-tagging tagged/partitioned nodes, the order in which we iterate the nodes will matter as each node will make a decision based on if the parent node is partitoned or has a compatible dtype. This change makes sure we iterate the nodes in the order of the original graph_module. Signed-off-by: Oscar Andersson Change-Id: I0e77c20440fee52f50cb1b7e57d8c39789b4f3c5 --- backends/arm/tosa_partitioner.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/backends/arm/tosa_partitioner.py b/backends/arm/tosa_partitioner.py index 738c5ab8204..de85dfae92f 100644 --- a/backends/arm/tosa_partitioner.py +++ b/backends/arm/tosa_partitioner.py @@ -96,7 +96,9 @@ def is_partitioned(node: torch.fx.Node, tag=tag) -> bool: # De-tag outmost q-nodes upwards and dq-nodes downwards. # De-tag if at least one input/ output is not part of partition. - for node in partition.nodes: + for node in exported_program.graph_module.graph.nodes: + if not is_partitioned(node): + continue if is_quant_node(node): for input in node.all_input_nodes: if not is_partitioned(input):