code: def batch_shuffle(x):
      # gather from all gpus
      batch_size_this = x.shape[0]
      all_xs, batch_size_all = concat_all_gather(x)
      all_xs_concat = torch.cat(all_xs, dim=0)
      total_bs = sum(batch_size_all)
  
      rank = dist.get_rank()
      assert batch_size_all[rank] == batch_size_this
  
      idx_range = (
          sum(batch_size_all[:rank]),
          sum(batch_size_all[: rank + 1]),
      )
  
      # random shuffle index
      idx_shuffle = torch.randperm(total_bs, device=x.device)
      # broadcast to all gpus
      dist.broadcast(idx_shuffle, src=0)
  
      # index for restoring
      idx_unshuffle = torch.argsort(idx_shuffle)
  
      # shuffled index for this gpu
      splits = torch.split(
          idx_shuffle, math.ceil(total_bs / dist.get_world_size())
      )
      if len(splits) > rank:
          idx_this = splits[rank]
      else:
          idx_this = idx_shuffle.new_zeros([0])
      return (
          all_xs_concat[idx_this],
          idx_unshuffle[idx_range[0] : idx_range[1]],
      )
  
file_path: detectron2/projects/Rethinking-BatchNorm/configs/mask_rcnn_BNhead_shuffle.py
imports: import math
  import torch
  import torch.distributed
  from detectron2.modeling.roi_heads import FastRCNNConvFCHead, MaskRCNNConvUpsampleHead
  from detectron2.utils import comm
  from fvcore.nn.distributed import differentiable_all_gather
  from .mask_rcnn_BNhead import model, dataloader, lr_multiplier, optimizer, train
label_desc: Use image node_id 2 for calling a global function with example usage: batch_shuffle(x) and returns: all_xs_concat,  idx_unshuffle