Skip to content

Commit 049a2f1

Browse files
committed
Check __cuda_stream__ attribute if method fails
1 parent 5376a4e commit 049a2f1

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

cuda_core/cuda/core/experimental/_stream.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,10 @@ def _init(obj=None, *, options: Optional[StreamOptions] = None):
8989
if obj is not None:
9090
if not hasattr(obj, "__cuda_stream__"):
9191
raise ValueError
92-
info = obj.__cuda_stream__()
92+
try:
93+
info = obj.__cuda_stream__()
94+
except TypeError:
95+
info = obj.__cuda_stream__
9396
assert info[0] == 0
9497
self._mnff.handle = cuda.CUstream(info[1])
9598
# TODO: check if obj is created under the current context/device

0 commit comments

Comments
 (0)