import torch


def to_sorted_tensor(tensor, lengths, sort_dim, device):
    """
    Sort tensor according to sequence lengths.
    This is used before applying pack_padded_sequence.
    :param tensor: input tensor
    :param lengths: 1D tensor of sequence lengths
    :param sort_dim: the dimension to sort in input tensor
    :param device: calculation device
    :return: sorted tensor, sorted lengths, and sorted index.
    """
    sorted_lengths, sorted_idx = torch.sort(
        lengths.long(), dim=0, descending=True)
    # sorted_idx = sorted_idx.to(device)
    sorted_tensor = tensor.index_select(dim=sort_dim, index=sorted_idx)
    return sorted_tensor, sorted_lengths, sorted_idx