Hi Nitin,
To send all tensors to one GPU, you’d want to use dist.gather, which will gather all of the tensors onto one gpu (this is assuming you have one process running per gpu). If your tensor is t, then your call would look like:
t = your_tensor(size)
if rank == 0:
# rank 0 is the node all tensors will be gathered on
    gathered_tensors = [torch.zeros(size) for _ in range(WORLD_SIZE)]
    dist.gather(t, gathered_tensors, dst=0)
else:
    dist.gather(t, dst=0)
    

Then, you can compute what you want and send the result back wrapped in a tensor via a dist.scatter. To make sure the tensors don’t change on the gpu before they are gathered, ensure all nodes call into gather at the same time when all nodes have the desired value. You could also use torch.distributed.barrier if you need additional synchronization. Check out the docs at https://pytorch.org/docs/stable/distributed.html 18