diff --git a/spacy/util.py b/spacy/util.py index 032c1741a..0d07a532d 100644 --- a/spacy/util.py +++ b/spacy/util.py @@ -317,13 +317,13 @@ def is_in_jupyter(): return False -def get_cuda_stream(require=False): +def get_cuda_stream(require=False, non_blocking=True): if CudaStream is None: return None elif isinstance(Model.ops, NumpyOps): return None else: - return CudaStream() + return CudaStream(non_blocking=non_blocking) def get_async(stream, numpy_array):