src/transformers/models/roberta/modeling_roberta.py/RobertaForCausalLM/_reorder_cache
class RobertaForCausalLM:    def _reorder_cache(self, past_key_values, beam_idx):
        reordered_past = ()
        for layer_past in past_key_values:
            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
        return reordered_past