Skip to content

Commit 186ac1b

Browse files
apboseperi044
authored andcommitted
backend changes to decomposition for transpose
1 parent b040d81 commit 186ac1b

File tree

1 file changed

+8
-7
lines changed

1 file changed

+8
-7
lines changed

py/torch_tensorrt/dynamo/backend/backends.py

+8-7
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from torch._dynamo.backends.common import aot_autograd
1111
from torch._dynamo.utils import detect_fake_mode
1212
from torch._functorch.aot_autograd import aot_export_joint_simple
13+
from torch._ops import OpOverload
1314
from torch_tensorrt.dynamo import CompilationSettings
1415
from torch_tensorrt.dynamo._compiler import compile_module
1516
from torch_tensorrt.dynamo.lowering import (
@@ -59,17 +60,17 @@ def aot_torch_tensorrt_aten_backend(
5960
_pretraced_backend, settings=settings, engine_cache=engine_cache
6061
)
6162
settings_aot_autograd = {}
62-
settings_aot_autograd["decompostions"] = get_decompositions(
63+
settings_aot_autograd["decompositions"] = get_decompositions(
6364
settings.enable_experimental_decompositions
6465
)
65-
# This is added since detach lowering leads to alias nodes
66-
# Error - View operation returned a tensor that is the same as the input base tensor
67-
# torch nop_decompositions in torch/_decomp/decompositions.py
68-
if aten.detach in settings_aot_autograd["decompositions"]:
69-
del settings_aot_autograd["decompositions"][aten.detach]
66+
# transpose key deleted since not desirable to lower it to permute
67+
for key in settings_aot_autograd["decompositions"]:
68+
if "transpose" in key._name:
69+
to_delete = key
70+
del settings_aot_autograd["decompositions"][to_delete]
7071
return aot_autograd(
7172
fw_compiler=_pretraced_backend_autograd,
72-
decompositions=get_decompositions(settings.enable_experimental_decompositions),
73+
decompositions=settings_aot_autograd["decompositions"],
7374
)(gm, sample_inputs)
7475

7576

0 commit comments

Comments
 (0)