src/transformers/models/transfo_xl/modeling_transfo_xl.py/TransfoXLLMHeadModel/prepare_inputs_for_generation
class TransfoXLLMHeadModel:    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **model_kwargs):
        inputs = {}

        # if past is defined in model kwargs then use it for faster decoding
        if past_key_values:
            inputs["mems"] = past_key_values
            inputs["input_ids"] = input_ids[:, -1].unsqueeze(-1)
        else:
            inputs["input_ids"] = input_ids

        return inputs