code: def batch_unshuffle_ddp(x, idx_unshuffle, rank, world_size):
      """
      Undo batch shuffle.
      *** Only support DistributedDataParallel (DDP) model. ***
      """
      # gather from all gpus
      batch_size_this = x.shape[0]
      x_gather = concat_all_gather(x)
      batch_size_all = x_gather.shape[0]
  
      # restored index for this gpu
      idx_this = idx_unshuffle.view(world_size, -1)[rank]
  
      return x_gather[idx_this]
  
label_desc: Use image node_id 3 for calling a global function with example usage: batch_unshuffle_ddp(x, idx_unshuffle, rank, world_size) and returns: x_gather