code: def batch_unshuffle(x, idx_unshuffle):
      all_x, _ = concat_all_gather(x)
      x_gather = torch.cat(all_x, dim=0)
      return x_gather[idx_unshuffle]
  
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 3 for calling a global function with example usage: batch_unshuffle(x, idx_unshuffle) and returns: x_gather