if we use torch.nn.parallel.gather to collect data from other GPUs, then we do some operations on the gathered data A, the gradient will go back to the original array( since it’s from the ‘nn’ package, that’s what supposed to happen).   We can check grad_fn of  A,  it says CopyBackwards, that’s what we expect.
but it seems the all_gather from torch.distributed package is different.  it simply copies the data here and doesn’t have any grad_fn function.  Does this mean if we use this function, we have to manually compute gradients for the gathered data?
related post:

  
    
    
    Something about dist.all_gather operation autograd
  
  
    I have a question. When we apply dist.all_gather() operation, suppose the there are 4 gpus, and each gpu will get the value of others, and when we apply the result of all_gather with ground truth to calculate loss, does loss can backward? or the dist.all_gather operation will break the graph like the operation of detach()?
  



  
    
    
    Is my code the correct way using DistributedDataParallel in single node multi GPUs? distributed
  
  
    I use torch.nn.parallel.DistributedDataParallel API in PyTorch1.1 to spawn my multi-card model (2 GPUs) to 8 GPUs.  According to official tutorial GETTING STARTED WITH DISTRIBUTED DATA PARALLEL, DistributedDataParallel is the recommanded way to parallel one’s model. I am not confident about my implementation and I can’t find other valuable tutorials, so come here for help. 
My ideas are simply as follow: 

split my 3D CNN model into 2 GPUs (simply called dev_in and dev_out),
use DistributedDataP…