Skip to content

Commit 07e16fd

Browse files
committed
fix(//py): Don't crash harshly on import when CUDA is not available
Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
1 parent 4a0a918 commit 07e16fd

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

Diff for: py/trtorch/Device.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,11 @@ def _from_torch_device(cls, torch_dev: torch.device):
107107

108108
@classmethod
109109
def _current_device(cls):
110-
dev = trtorch._C._get_current_device()
110+
try:
111+
dev = trtorch._C._get_current_device()
112+
except RuntimeError:
113+
trtorch.logging.log(trtorch.logging.Level.Error, "Cannot get current device")
114+
return None
111115
return cls(gpu_id=dev.gpu_id)
112116

113117
@staticmethod

0 commit comments

Comments
 (0)