def sequence_to_segments(
    sequence: Tensor,
    time_axis: int,
    lengths: Tensor,
    segment_size: Optional[int] = None,
    extra_left_context: int = 0,
    extra_right_context: int = 0,
) -> List[Tuple[Tensor, Tensor]]:
    """Breaks sequence into segments."""

    sequence = pad_sequence(
        sequence=sequence,
        time_axis=time_axis,
        extra_left_context=extra_left_context,
        extra_right_context=extra_right_context,
    )

    lengths = lengths + extra_left_context + extra_right_context

    segments: List[Tuple[Tensor, Tensor]] = []

    if segment_size is None:
        segments.append((sequence, lengths))
        return segments

    offset = 0
    end = sequence.shape[time_axis]
    step = segment_size
    size = extra_left_context + segment_size + extra_right_context

    while offset + extra_left_context + extra_right_context < end:
        clamped_size = min(size, end - offset)
        segment_lengths = torch.clamp(lengths - offset, min=0, max=clamped_size)
        indices = torch.arange(
            start=offset,
            end=(offset + clamped_size),
            step=1,
            dtype=torch.long,
            device=sequence.device,
        )
        segment_tensor = torch.index_select(sequence, time_axis, indices)
        segments.append((segment_tensor, segment_lengths))
        offset = offset + step

    return segments
Breaks sequence into segments.
fairseq/models/speech_to_text/utils.py