|
53 | 53 | logger = logging.getLogger(__name__)
|
54 | 54 |
|
55 | 55 |
|
56 |
| -@needs_cross_compile |
| 56 | +@needs_cross_compile # type: ignore |
57 | 57 | def cross_compile_for_windows(
|
58 | 58 | exported_program: ExportedProgram,
|
59 | 59 | inputs: Optional[Sequence[Sequence[Any]]] = None,
|
@@ -141,7 +141,7 @@ def cross_compile_for_windows(
|
141 | 141 | assume_dynamic_shape_support (bool): Setting this to true enables the converters work for both dynamic and static shapes. Default: False
|
142 | 142 | sparse_weights (bool): Enable sparsity for convolution and fully connected layers.
|
143 | 143 | enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels
|
144 |
| - capability (torch_tensorrt.EngineCapability): Restrict kernel selection to safe gpu kernels or safe dla kernels |
| 144 | + engine_capability (torch_tensorrt.EngineCapability): Restrict kernel selection to safe gpu kernels or safe dla kernels |
145 | 145 | num_avg_timing_iters (int): Number of averaging timing iterations used to select kernels
|
146 | 146 | workspace_size (int): Maximum size of workspace given to TensorRT
|
147 | 147 | dla_sram_size (int): Fast software managed RAM used by DLA to communicate within a layer.
|
@@ -479,7 +479,7 @@ def compile(
|
479 | 479 | assume_dynamic_shape_support (bool): Setting this to true enables the converters work for both dynamic and static shapes. Default: False
|
480 | 480 | sparse_weights (bool): Enable sparsity for convolution and fully connected layers.
|
481 | 481 | enabled_precision (Set(Union(torch.dtype, torch_tensorrt.dtype))): The set of datatypes that TensorRT can use when selecting kernels
|
482 |
| - capability (torch_tensorrt.EngineCapability): Restrict kernel selection to safe gpu kernels or safe dla kernels |
| 482 | + engine_capability (torch_tensorrt.EngineCapability): Restrict kernel selection to safe gpu kernels or safe dla kernels |
483 | 483 | num_avg_timing_iters (int): Number of averaging timing iterations used to select kernels
|
484 | 484 | workspace_size (int): Maximum size of workspace given to TensorRT
|
485 | 485 | dla_sram_size (int): Fast software managed RAM used by DLA to communicate within a layer.
|
@@ -723,7 +723,7 @@ def compile(
|
723 | 723 | return trt_gm
|
724 | 724 |
|
725 | 725 |
|
726 |
| -@fn_supports_debugger |
| 726 | +@fn_supports_debugger # type: ignore |
727 | 727 | def compile_module(
|
728 | 728 | gm: torch.fx.GraphModule,
|
729 | 729 | sample_arg_inputs: Sequence[Input],
|
@@ -1289,7 +1289,7 @@ def convert_exported_program_to_serialized_trt_engine(
|
1289 | 1289 | return serialized_engine
|
1290 | 1290 |
|
1291 | 1291 |
|
1292 |
| -@needs_cross_compile |
| 1292 | +@needs_cross_compile # type: ignore |
1293 | 1293 | def save_cross_compiled_exported_program(
|
1294 | 1294 | gm: torch.fx.GraphModule,
|
1295 | 1295 | file_path: str,
|
|
0 commit comments