def segments_to_sequence(
    segments: List[Tuple[Tensor, Tensor]], time_axis: int
) -> Tuple[Tensor, Tensor]:
    """Concatenate segments into a full sequence."""
    if len(segments) == 1:
        return segments[0]

    tensors_to_concat: List[Tensor] = []
    lengths_to_stack: List[Tensor] = []

    for tensor, lengths in segments:
        tensors_to_concat.append(tensor)
        lengths_to_stack.append(lengths)

    sequence = torch.cat(tensors_to_concat, dim=time_axis)
    lengths = torch.stack(lengths_to_stack, dim=0)
    lengths = torch.sum(lengths, dim=0)

    return sequence, lengths
Concatenate segments into a full sequence.
fairseq/models/speech_to_text/utils.py