Skip to content

Commit 9ea829d

Browse files
committed
chore: address review comments
Signed-off-by: Dheeraj Peri <[email protected]>
1 parent fbe929f commit 9ea829d

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

py/torch_tensorrt/dynamo/aten_tracer.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,10 @@
77
import torch
88
from torch._export import dynamic_dim, export
99
from torch_tensorrt._Input import Input
10-
from torch_tensorrt.dynamo._defaults import default_device
10+
from torch_tensorrt.dynamo._defaults import (
11+
ENABLE_EXPERIMENTAL_DECOMPOSITIONS,
12+
default_device,
13+
)
1114
from torch_tensorrt.dynamo.lowering import get_decompositions
1215
from torch_tensorrt.dynamo.utils import get_torch_inputs, set_log_level, to_torch_device
1316

@@ -75,7 +78,7 @@ def trace(
7578
trace_inputs.append(torch_inputs[idx])
7679

7780
experimental_decompositions = kwargs.get(
78-
"enable_experimental_decompositions", False
81+
"enable_experimental_decompositions", ENABLE_EXPERIMENTAL_DECOMPOSITIONS
7982
)
8083
with unittest.mock.patch(
8184
"torch._export.DECOMP_TABLE", get_decompositions(experimental_decompositions)

0 commit comments

Comments
 (0)