code: def batch_unshuffle_distributed(
      batch: torch.Tensor, shuffle: torch.Tensor
  ) -> torch.Tensor:
      """Undo batch shuffle over multiple devices.
  
      This code was taken and adapted from here:
      https://github.com/facebookresearch/moco.
  
      Args:
          batch:
              The tensor to unshuffle.
          shuffle:
              Index to restore the original tensor.
  
      Returns:
          The unshuffled tensor.
  
      """
      # gather from all devices
      batch_size_this = batch.shape[0]
      batch_gather = concat_all_gather(batch)
      batch_size_all = batch_gather.shape[0]
  
      num_devices = batch_size_all // batch_size_this
  
      # restored index for this gpu
      rank = dist.get_rank()
      idx_this = shuffle.view(num_devices, -1)[rank]
      return batch_gather[idx_this]
  
imports: import math
  import warnings
  from typing import Iterable, List, Optional, Tuple, Union
  import numpy
  import torch
  import torch.distributed
  import torch.nn
  from numpy.typing import NDArray
  from torch.nn import Module
  from torch.nn.modules.batchnorm import _BatchNorm
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 5 for calling a global function with example usage: batch_unshuffle_distributed(batch, shuffle) and returns: batch_gather