Skip to content

Commit 7993bb2

Browse files
Arm backend: Update partitioner de-tagging iteration order (#10813)
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 <[email protected]>
1 parent b30f912 commit 7993bb2

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

backends/arm/tosa_partitioner.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,9 @@ def is_partitioned(node: torch.fx.Node, tag=tag) -> bool:
9696

9797
# De-tag outmost q-nodes upwards and dq-nodes downwards.
9898
# De-tag if at least one input/ output is not part of partition.
99-
for node in partition.nodes:
99+
for node in exported_program.graph_module.graph.nodes:
100+
if not is_partitioned(node):
101+
continue
100102
if is_quant_node(node):
101103
for input in node.all_input_nodes:
102104
if not is_partitioned(input):

0 commit comments

Comments
 (0)