code: def all_gather_heterogeneous(
      sizes: List[int], x: torch.Tensor
  ) -> List[torch.Tensor]:
      """
      Gather a list of heterogeenous tensors shape in the first
      dimension (different batch sizes)
      """
      current_device = torch.device("cuda", torch.cuda.current_device())
      shape = x.shape[1:]
      all_x = [
          torch.zeros(
              size=(sizes[i], *shape),
              device=current_device,
              dtype=x.dtype,
          )
          for i in range(torch.distributed.get_world_size())
      ]
      torch.distributed.all_gather(all_x, x)
      return all_x
  
file_path: vissl/vissl/utils/distributed_utils.py
imports: from typing import List
  import torch
  import torch.distributed
  from classy_vision.generic.distributed_util import convert_to_distributed_tensor, convert_to_normal_tensor, is_distributed_training_run
label_desc: Use image node_id 3 for calling a global function with example usage: all_gather_heterogeneous(sizes, x) and returns: all_x