From 07e16fd7d7584d91adef1c698dd80e46200ac0f7 Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Wed, 20 Oct 2021 10:45:21 -0700 Subject: [PATCH] fix(//py): Don't crash harshly on import when CUDA is not available Signed-off-by: Naren Dasan Signed-off-by: Naren Dasan --- py/trtorch/Device.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/py/trtorch/Device.py b/py/trtorch/Device.py index 41a1308518..4a9395415e 100644 --- a/py/trtorch/Device.py +++ b/py/trtorch/Device.py @@ -107,7 +107,11 @@ def _from_torch_device(cls, torch_dev: torch.device): @classmethod def _current_device(cls): - dev = trtorch._C._get_current_device() + try: + dev = trtorch._C._get_current_device() + except RuntimeError: + trtorch.logging.log(trtorch.logging.Level.Error, "Cannot get current device") + return None return cls(gpu_id=dev.gpu_id) @staticmethod