code: def batch_shuffle_ddp(x, rank, world_size):
      """
      Batch shuffle, for making use of BatchNorm.
      *** 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]
  
      # random shuffle index
      idx_shuffle = torch.randperm(batch_size_all).cuda()
  
      # broadcast to all gpus
      torch.distributed.broadcast(idx_shuffle, src=0)
  
      # index for restoring
      idx_unshuffle = torch.argsort(idx_shuffle)
  
      # shuffled index for this gpu
      idx_this = idx_shuffle.view(world_size, -1)[rank]
  
      return x_gather[idx_this], idx_unshuffle
  
label_desc: Use image node_id 2 for calling a global function with example usage: batch_shuffle_ddp(x, rank, world_size) and returns: x_gather,  idx_unshuffle