Skip to content

Commit 776839a

Browse files
committed
Update device logging and defaults
1 parent 431b145 commit 776839a

File tree

2 files changed

+4
-7
lines changed

2 files changed

+4
-7
lines changed

py/torch_tensorrt/dynamo/_defaults.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,4 @@
1212
TRUNCATE_LONG_AND_DOUBLE = False
1313
USE_PYTHON_RUNTIME = False
1414
USE_FAST_PARTITIONER = True
15-
DEVICE = Device._current_device()
15+
DEVICE = Device(gpu_id=torch.cuda.current_device())

py/torch_tensorrt/dynamo/utils.py

+3-6
Original file line numberDiff line numberDiff line change
@@ -182,14 +182,11 @@ def parse_dynamo_kwargs(kwargs: Any) -> CompilationSettings:
182182
settings.device = to_torch_tensorrt_device(settings.device)
183183

184184
# Check and update device settings
185-
default_torch_gpu_idx = torch.cuda.default_stream().device.index
186-
if "device" not in kwargs and default_torch_gpu_idx != settings.device.gpu_id:
187-
logger.warning(
188-
f"No device specified, detected differing gpu IDs for CUDA default: {settings.device.gpu_id} "
189-
f"and Torch default: {default_torch_gpu_idx}. Using Torch default gpu ID: {default_torch_gpu_idx}. "
185+
if "device" not in kwargs:
186+
logger.info(
187+
f"Device not specified, using Torch default current device - cuda:{settings.device.gpu_id}. "
190188
"If this is incorrect, please specify an input device, via the device keyword."
191189
)
192-
settings.device = Device(gpu_id=default_torch_gpu_idx)
193190

194191
logger.debug(f"Compiling with Settings:\n{settings}")
195192

0 commit comments

Comments
 (0)