src/transformers/models/gpt2/modeling_tf_gpt2.py/TFGPT2LMHeadModel/prepare_inputs_for_generation
class TFGPT2LMHeadModel:    def prepare_inputs_for_generation(self, inputs, past, **kwargs):
        # only last token for inputs_ids if past is defined in kwargs
        if past:
            inputs = tf.expand_dims(inputs[:, -1], -1)

        return {"input_ids": inputs, "past": past, "use_cache": kwargs["use_cache"]}