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]>
pytorch · Oct 20, 2021 · 07e16fd · 07e16fd
1 parent 4a0a918
commit 07e16fd
Showing 1 changed file with 5 additions and 1 deletion.
diff --git 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