import torch


def reorder_gru_states(gru_state, order):
    """
    gru_state: [num_layer * num_directions, batch, hidden_size]
    order: list of sequence length
    """
    assert len(order) == gru_state.size(1)
    order = torch.LongTensor(order).to(gru_state.device)
    sorted_state = gru_state.index_select(index=order, dim=1)
    return sorted_state