How to assign a specific stream for ProcessGroupNCCL and subsequent communication operations?

1

I want to pass an externally allocated stream to NCCL in PyTorch so that the related communication operations are executed on that stream.

Thank you for all the replies!