Commit 29b7d374 authored by Xiaochuan Ye's avatar Xiaochuan Ye Committed by LeiWang1999
Browse files

[Bugfix]Add CUDA availability check in CtypesKernelAdapter (#267)

* fix: Add CUDA availability check in CtypesKernelAdapter

* fix: Add CUDA availability check in CythonKernelWrapper
parent 5f5bf53c
...@@ -226,7 +226,10 @@ class CtypesKernelAdapter(BaseKernelAdapter): ...@@ -226,7 +226,10 @@ class CtypesKernelAdapter(BaseKernelAdapter):
# if stream is not None, we need to pass the stream to the library # if stream is not None, we need to pass the stream to the library
if stream is None: if stream is None:
stream = torch.cuda.current_stream().cuda_stream if str(self.target).startswith("cuda") and torch.cuda.is_available():
stream = torch.cuda.current_stream().cuda_stream
else:
stream = 0
self._forward_from_prebuild_lib(*args, stream=stream) self._forward_from_prebuild_lib(*args, stream=stream)
......
...@@ -73,7 +73,10 @@ cdef class CythonKernelWrapper: ...@@ -73,7 +73,10 @@ cdef class CythonKernelWrapper:
# Use current CUDA stream if none specified # Use current CUDA stream if none specified
if stream == -1: if stream == -1:
stream = torch.cuda.current_stream().cuda_stream if torch.cuda.is_available():
stream = torch.cuda.current_stream().cuda_stream
else:
stream = 0
cdef int ins_idx = 0 cdef int ins_idx = 0
cdef list tensor_list = [] cdef list tensor_list = []
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment