pytorch

b82453cb - Run dist_autograd backward RPCs on appropriate CUDA streams. (#60606)

Run dist_autograd backward RPCs on appropriate CUDA streams. (#60606) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60606 TensorPipe receives tensors over the wire on custom streams and these streams are passed to some RPC callbacks but not to `BACKWARD_AUTOGRAD_REQ`. As a result, `BACKWARD_AUTOGRAD_REQ` ran on the default stream while still using tensors from the custom stream. This resulted in downstream autograd operations running on the incorrect stream. To fix this, I've passed the streams to `BACKWARD_AUTOGRAD_REQ` as well and added an appropriate guard. #Closes: https://github.com/pytorch/pytorch/issues/59793 ghstack-source-id: 132252069 Test Plan: Test https://github.com/pytorch/pytorch/issues/59793 Reviewed By: mrshenli Differential Revision: D29347244 fbshipit-source-id: 8ff8b150763c970ab15c2cac8dccf56e66e9ef5d