Check status after calls to cudaGetErrorName/String

rapidsai · Jan 13, 2022 · 8be6c0b · 8be6c0b
1 parent 84994cb
commit 8be6c0b
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/python/rmm/_cuda/gpu.py b/python/rmm/_cuda/gpu.py
@@ -6,8 +6,15 @@
 class CUDARuntimeError(RuntimeError):
     def __init__(self, status: cudart.cudaError_t):
         self.status = status
-        _, name = cudart.cudaGetErrorName(status)
-        _, msg = cudart.cudaGetErrorString(status)
+
+        err, name = cudart.cudaGetErrorName(status)
+        if err != cudart.cudaError_t.cudaSuccess:
+            raise CUDARuntimeError(err.value)
+
+        err, msg = cudart.cudaGetErrorString(status)
+        if err != cudart.cudaError_t.cudaSuccess:
+            raise CUDARuntimeError(err.value)
+
         super(CUDARuntimeError, self).__init__(
             "%s: %s" % (name.decode(), msg.decode())
         )