docs: def get_world_size():
      """
      Gets total number of distributed workers or returns one if distributed is
      not initialized.
      """
      if torch.distributed.is_available() and torch.distributed.is_initialized():
          world_size = torch.distributed.get_world_size()
      else:
          world_size = 1
      return world_size