Flips a padded tensor along the time dimension without affecting masked entries.

        Parameters
        ----------
        padded_sequence : ``torch.Tensor``
            The tensor to flip along the time dimension.
            Assumed to be of dimensions (batch size, num timesteps, ...)
        sequence_lengths : ``torch.Tensor``
            A list containing the lengths of each unpadded sequence in the batch.

        Returns
        -------
        A ``torch.Tensor`` of the same shape as padded_sequence.

def masked_flip(padded_sequence: torch.Tensor,
                sequence_lengths: List[int]) -> torch.Tensor:
    """
        Flips a padded tensor along the time dimension without affecting masked entries.

        Parameters
        ----------
        padded_sequence : ``torch.Tensor``
            The tensor to flip along the time dimension.
            Assumed to be of dimensions (batch size, num timesteps, ...)
        sequence_lengths : ``torch.Tensor``
            A list containing the lengths of each unpadded sequence in the batch.

        Returns
        -------
        A ``torch.Tensor`` of the same shape as padded_sequence.
        """
    assert padded_sequence.size(0) == len(sequence_lengths), \
        f'sequence_lengths length ${len(sequence_lengths)} does not match batch size ${padded_sequence.size(0)}'
    num_timesteps = padded_sequence.size(1)
    flipped_padded_sequence = torch.flip(padded_sequence, [1])
    sequences = [flipped_padded_sequence[i, num_timesteps - length:] for i, length in enumerate(sequence_lengths)]
    return torch.nn.utils.rnn.pad_sequence(sequences, batch_first=True)