examples/research_projects/bertabs/modeling_bertabs.py/TransformerDecoderState/update_state
class TransformerDecoderState:    def update_state(self, new_input, previous_layer_inputs):
        state = TransformerDecoderState(self.src)
        state.previous_input = new_input
        state.previous_layer_inputs = previous_layer_inputs
        return state