Skip to content

Commit 808e603

Browse files
authored
Merge pull request #679 from NVIDIA/python_packaging_feats
fix(//py): Don't crash harshly on import when CUDA is not available
2 parents b683b6e + 07e16fd commit 808e603

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

py/trtorch/Device.py

Lines changed: 5 additions & 1 deletion
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)