code: def check_and_append_tensor_for_gather(
      group, rank, world_size, input_
  ):
      # gather the size of the first dimension of the tensor from all ranks
      current_length = input_.size()[0]
      first_dim = torch.tensor(
          [[current_length]], device=torch.cuda.current_device()
      )
      input_list = [
          torch.empty_like(first_dim) for _ in range(world_size)
      ]
      input_list[rank].copy_(first_dim)
      torch.distributed.all_gather(input_list, first_dim, group=group)
      all_input_list = torch.cat(input_list, dim=0).contiguous()
      max_length = torch.max(all_input_list)
  
      # if the size are different than the max, extend the tensor
      # accordingly
      if max_length > current_length:
          padding = tuple([0] * (input_.dim() * 2 - 1)) + tuple(
              [max_length - current_length]
          )
          input_ = F.pad(input=input_, pad=padding)
  
      return input_
  
file_path: megatron-lm/tasks/orqa/supervised/finetune.py
imports: from functools import partial
  import sys
  import math
  import torch
  import torch.nn.functional
  from megatron import get_args, get_timers, get_tokenizer, print_rank_0
  from megatron.core import mpu
  from megatron.indexer import IndexBuilder
  from megatron.model.biencoder_model import biencoder_model_provider
  from megatron.utils import average_losses_across_data_parallel_group
  from pretrain_ict import get_group_world_size_rank
  from tasks.finetune_utils import finetune
  from tasks.orqa.supervised.eval_utils import accuracy_func_provider
  from tasks.orqa.supervised.eval_utils import process_batch, task_collate_fn
  from tasks.orqa.evaluate_utils import ORQAEvaluator
label_desc: Use image node_id 1 for calling a global function with example usage: check_and_append_tensor_for_gather(group, rank, world_size, input_) and returns: input_