def reorder_attn_buffer(input_buffer, new_order):
    """Reorder buffered internal state (for incremental generation)."""
    # input_buffer = self._get_input_buffer(incremental_state)
    for k in input_buffer.keys():
        input_buffer_k = input_buffer[k]
        if input_buffer_k is not None:
            input_buffer[k] = input_buffer_k.index_select(0, new_order)
        # incremental_state = self._set_input_buffer(incremental_state, input_buffer)
    return input_buffer
Reorder buffered internal state (for incremental generation).
xtune/src/transformers/modeling_bart.py