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