Skip to content

Commit a226d84

Browse files
Merge branch 'main' into change-989421
2 parents 6a48cab + 9746ce7 commit a226d84

File tree

1 file changed

+1
-2
lines changed

1 file changed

+1
-2
lines changed

backends/arm/_passes/decompose_select.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,14 +37,13 @@ def call(self, graph_module: torch.fx.GraphModule):
3737
rank = len(input_node.meta["val"].size())
3838
dim = dim % rank if dim < 0 else dim
3939
index = index % rank if index < 0 else index
40-
dim_list = list(range(rank))
4140

4241
with graph_module.graph.inserting_before(node):
4342
slice_node = create_node(
4443
graph_module.graph, slice_op, (input_node, dim, index, index + 1)
4544
)
4645
squeeze_node = create_node(
47-
graph_module.graph, squeeze_op, (slice_node, dim_list)
46+
graph_module.graph, squeeze_op, (slice_node, [dim])
4847
)
4948

5049
node.replace_all_uses_with(squeeze_node)

0 commit comments

Comments
 (0)