def pad_batch(batch, pad_id, neox_args):
    """
    pads context lengths in batch with pad_id to equal neox_args.seq_length,
    and returns the padded batch and the new lengths.

    batch: torch.Tensor of tokens
    pad_id: int, integer to use as padding token
    neox_args: neox_args
    """

    context_lengths = []
    for tokens in batch:
        context_length = len(tokens)
        if context_length < neox_args.seq_length:
            tokens.extend([pad_id] * (neox_args.seq_length - context_length))
        context_lengths.append(context_length)
    return batch, context_lengths