|
3 | 3 | from typing import Any, Callable, Dict, Optional, Sequence |
4 | 4 |
|
5 | 5 | import torch |
| 6 | +import torch_tensorrt |
6 | 7 | from torch_tensorrt._Device import Device |
7 | 8 | from torch_tensorrt._Input import Input |
8 | 9 | from torch_tensorrt.dynamo import CompilationSettings |
| 10 | +from torch_tensorrt.dynamo._defaults import PRECISION |
9 | 11 |
|
10 | 12 | from packaging import version |
11 | 13 |
|
@@ -159,6 +161,28 @@ def parse_dynamo_kwargs(kwargs: Any) -> CompilationSettings: |
159 | 161 | if settings.debug: |
160 | 162 | logger.setLevel(logging.DEBUG) |
161 | 163 |
|
| 164 | + # TODO: Remove once Dynamo precisions refactoring is complete |
| 165 | + if "enabled_precisions" in kwargs: |
| 166 | + enabled_precisions = kwargs["enabled_precisions"] |
| 167 | + |
| 168 | + if ( |
| 169 | + torch.float16 in enabled_precisions |
| 170 | + or torch_tensorrt.dtype.half in enabled_precisions |
| 171 | + ): |
| 172 | + settings.precision = torch.float16 |
| 173 | + elif ( |
| 174 | + torch.float32 in enabled_precisions |
| 175 | + or torch_tensorrt.dtype.float in enabled_precisions |
| 176 | + ): |
| 177 | + settings.precision = torch.float32 |
| 178 | + elif len(enabled_precisions) == 0: |
| 179 | + logger.info(f"No precision specified, defaulting to {PRECISION}") |
| 180 | + settings.precision = PRECISION |
| 181 | + else: |
| 182 | + raise ValueError( |
| 183 | + f"Precision {enabled_precisions} not supported in the Dynamo Path" |
| 184 | + ) |
| 185 | + |
162 | 186 | # Parse input runtime specification |
163 | 187 | settings.use_python_runtime = use_python_runtime_parser(settings.use_python_runtime) |
164 | 188 |
|
|
0 commit comments