We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent dfc4668 commit 01bf54eCopy full SHA for 01bf54e
bitsandbytes/functional.py
@@ -441,8 +441,9 @@ def is_on_gpu(tensors: Iterable[torch.Tensor]):
441
return on_gpu
442
443
444
-def get_tensor_stream(tensor: Tensor) -> torch.cuda.Stream:
445
- return torch.cuda.current_stream(tensor.device)
+def get_tensor_stream(tensor: Tensor) -> ct.c_void_p:
+ # We use the raw stream for performance reasons.
446
+ return ct.c_void_p(torch._C._cuda_getCurrentRawStream(tensor.device.index))
447
448
449
def get_ptr(A: Optional[Tensor]) -> Optional[ct.c_void_p]:
0 commit comments