@leimao
No, I mean, you can use libtorch exposed api to get CUDA stream.
I wrote this recently ‘https://pytorch.org/cppdocs/notes/tensor_cuda_stream.html 4’ , not finished yet, someone said it is not working well, but no follow up, can you try and let me know if it works? The other thing is that please pay attention to setCurrentStream if you plan to use it, very tricky.