File tree Expand file tree Collapse file tree 1 file changed +5
-1
lines changed Expand file tree Collapse file tree 1 file changed +5
-1
lines changed Original file line number Diff line number Diff line change 4
4
5
5
import torch
6
6
import torch_tensorrt ._C .ts as _C
7
- from torch_tensorrt import _enums
8
7
from torch_tensorrt ._Device import Device
9
8
from torch_tensorrt ._Input import Input
10
9
from torch_tensorrt .ts ._compile_spec import _parse_compile_spec , _parse_device
11
10
11
+ from torch_tensorrt import _enums
12
+
12
13
13
14
def compile (
14
15
module : torch .jit .ScriptModule ,
@@ -137,6 +138,9 @@ def compile(
137
138
"capability" : capability , # Restrict kernel selection to safe gpu kernels or safe dla kernels
138
139
"num_avg_timing_iters" : num_avg_timing_iters , # Number of averaging timing iterations used to select kernels
139
140
"workspace_size" : workspace_size , # Maximum size of workspace given to TensorRT
141
+ "dla_sram_size" : dla_sram_size ,
142
+ "dla_local_dram_size" : dla_local_dram_size ,
143
+ "dla_global_dram_size" : dla_global_dram_size ,
140
144
"calibrator" : calibrator ,
141
145
"truncate_long_and_double" : truncate_long_and_double ,
142
146
"torch_fallback" : {
You can’t perform that action at this time.
0 commit comments