Simplify handling of Numba streams

rapidsai · Dec 11, 2020 · 079d24d · 079d24d
1 parent fadf308
commit 079d24d
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/python/rmm/_cuda/stream.pyx b/python/rmm/_cuda/stream.pyx
@@ -94,10 +94,7 @@ cdef class Stream:
         return self.c_is_default()
 
     def _init_from_numba_stream(self, obj):
-        if obj.handle.value is None:
-            self._cuda_stream = cuda_stream_default.value()
-        else:
-            self._cuda_stream = <cudaStream_t><uintptr_t>(obj.handle.value)
+        self._cuda_stream = <cudaStream_t><uintptr_t>(int(obj))
         self._owner = obj
 
     def _init_from_cupy_stream(self, obj):