Skip to content

Commit 8e3de23

Browse files
committed
rebase
1 parent 8cedc03 commit 8e3de23

File tree

1 file changed

+15
-13
lines changed

1 file changed

+15
-13
lines changed

py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py

+15-13
Original file line numberDiff line numberDiff line change
@@ -17,19 +17,21 @@
1717
from .repair_input_as_output import repair_input_as_output
1818
from .replace_max_pool_with_indices import replace_max_pool_with_indices
1919

20-
ATEN_POST_LOWERING_PASSES = DynamoPassManager.build_from_passlist(
21-
[
22-
remove_input_alias_fixing_clones,
23-
constant_fold,
24-
repair_input_as_output,
25-
fuse_prims_broadcast,
26-
fuse_distributed_ops,
27-
replace_max_pool_with_indices,
28-
remove_assert_nodes,
29-
accumulate_fp32_matmul,
30-
remove_num_users_is_0_nodes,
31-
]
32-
)
20+
pass_list = [
21+
remove_input_alias_fixing_clones,
22+
constant_fold,
23+
repair_input_as_output,
24+
fuse_prims_broadcast,
25+
replace_max_pool_with_indices,
26+
remove_assert_nodes,
27+
accumulate_fp32_matmul,
28+
remove_num_users_is_0_nodes,
29+
]
30+
31+
if not is_tegra_platform():
32+
pass_list.append(fuse_distributed_ops)
33+
34+
ATEN_POST_LOWERING_PASSES = DynamoPassManager.build_from_passlist(pass_list)
3335

3436
ATEN_PRE_LOWERING_PASSES = DynamoPassManager.build_from_passlist(
3537
[

0 commit comments

Comments
 (0)