import torch


def to_original_tensor(sorted_tensor, sorted_idx, sort_dim, device):
    """
    Restore tensor to its original order.
    This is used after applying pad_packed_sequence.
    :param sorted_tensor: a sorted tensor
    :param sorted_idx: sorted index of the sorted_tensor
    :param sort_dim: the dimension of sorted_tensor where it is sorted
    :device: calculation device
    :return: the original unsorted tensor
    """
    original_idx = torch.LongTensor(sorted(
        range(sorted_idx.size(0)), key=lambda x: sorted_idx[x])).to(device)
    tensor = sorted_tensor.index_select(dim=sort_dim, index=original_idx)
    return tensor