code: def all_gather_sizes(x: torch.Tensor) -> List[int]:
      """
      Get the first dimension sizes of the the tensor to gather on each
      of the distributed workers
      """
      dist_rank = torch.distributed.get_rank()
      world_size = torch.distributed.get_world_size()
      current_device = torch.device("cuda", torch.cuda.current_device())
      sizes = torch.zeros(
          size=(world_size,), device=current_device, dtype=torch.int64
      )
      sizes[dist_rank] = x.shape[0]
      torch.distributed.all_reduce(sizes)
      return list(sizes.cpu().numpy())
  
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 2 for calling a global function with example usage: all_gather_sizes(x) and returns: list