Synchronize a specific cuda stream in python

In PyTorch, we can use torch.cuda.synchronize to synchronize all streams on a CUDA device. Is there any way to synchronize a specific stream in python like cudaStreamSynchronize from C++ runtime API?

You could use torch.cuda.Stream.synchronize() or torch.cuda.Stream.wait_stream(stream) to synchronize with the passed CUDA stream.

1 Like