File tree Expand file tree Collapse file tree 1 file changed +5
-2
lines changed Expand file tree Collapse file tree 1 file changed +5
-2
lines changed Original file line number Diff line number Diff line change 7
7
import torch
8
8
from torch ._export import dynamic_dim , export
9
9
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
+ )
11
14
from torch_tensorrt .dynamo .lowering import get_decompositions
12
15
from torch_tensorrt .dynamo .utils import get_torch_inputs , set_log_level , to_torch_device
13
16
@@ -75,7 +78,7 @@ def trace(
75
78
trace_inputs .append (torch_inputs [idx ])
76
79
77
80
experimental_decompositions = kwargs .get (
78
- "enable_experimental_decompositions" , False
81
+ "enable_experimental_decompositions" , ENABLE_EXPERIMENTAL_DECOMPOSITIONS
79
82
)
80
83
with unittest .mock .patch (
81
84
"torch._export.DECOMP_TABLE" , get_decompositions (experimental_decompositions )
You can’t perform that action at this time.
0 commit comments