diff --git a/backends/arm/_passes/annotate_channels_last_dim_order_pass.py b/backends/arm/_passes/annotate_channels_last_dim_order_pass.py index cfe9340c3b0..ce15d8298c3 100644 --- a/backends/arm/_passes/annotate_channels_last_dim_order_pass.py +++ b/backends/arm/_passes/annotate_channels_last_dim_order_pass.py @@ -129,6 +129,7 @@ def insert_input_transpose(node, input_node, graph_module): permute_node.meta["tosa_dim_order"] = tuple( range(len(input_node.meta["val"].size())) ) + permute_node.meta["val"] = input_node.meta["val"] @staticmethod def insert_output_transpose(node, graph_module): @@ -141,6 +142,9 @@ def insert_output_transpose(node, graph_module): permute_node.meta["tosa_dim_order"] = ( AnnotateChannelsLastDimOrder.NHWC_order ) + permute_node.meta["val"] = node.meta["val"].permute( + AnnotateChannelsLastDimOrder.NHWC_order + ) node.meta["tosa_dim_order"] = (0, 1, 2, 3) users = [user for user in node.users if user != permute_node] for user in users: