import torch


def sort_sequence(tensor, lengths, batch_first=False):
    """
    Sorts sequence in descending order
    tensor: Padded tensor of variable length stuff (Torch tensor)
    lengths: Lengths of padded tensor (Torch LongTensor)
    batch_first: Boolean, whether tensor is batch_first or not  
    """
    idx = None
    if batch_first:
        idx = 0
    else:
        idx = 1 

    sorted_lengths, indices = torch.sort(lengths, dim=0, descending=True)
    new_tensor = torch.index_select(tensor, idx, indices)
    return new_tensor, sorted_lengths, indices