There is no way to do this today. It also depends on which distributed backend you’re using whether this would be possible in the first place. If you’re using Gloo, it might be possible, but it’s quite a bit of work. If you’re using NCCL, it’s up to NVIDIA. If you’re using MPI, I don’t know.