provider_options = {
    "trt_engine_cache_enable": True,
    "trt_engine_cache_path": "tmp/trt_cache_gpt2_example",
    "trt_profile_min_shapes": "input_ids:1x16,attention_mask:1x16",
    "trt_profile_max_shapes": "input_ids:1x64,attention_mask:1x64",
    "trt_profile_opt_shapes": "input_ids:1x64,attention_mask:1x64",
}

ort_model = ORTModelForCausalLM.from_pretrained(
    "gpt2",
    export=True,
    use_cache=False,
    provider="TensorrtExecutionProvider",
    provider_options=provider_options,
)

ort_model.generate(
    input_ids=torch.tensor([[1] * 16]).to("cuda"),
    max_new_tokens=64-16,
    min_new_tokens=64-16,
    pad_token_id=0,
    eos_token_id=0,
)