add _supports_cache_class and _supports_static_cache

huggingface · Jun 29, 2024 · fba9586 · fba9586
1 parent aa25c34
commit fba9586
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/optimum/onnxruntime/modeling_decoder.py b/optimum/onnxruntime/modeling_decoder.py
@@ -122,6 +122,9 @@ class ORTModelForCausalLM(ORTModel, GenerationMixin):
     auto_model_class = AutoModelForCausalLM
     main_input_name = "input_ids"
 
+    _supports_cache_class = False
+    _supports_static_cache = False
+
     def __init__(
         self,
         model: onnxruntime.InferenceSession,

diff --git a/optimum/onnxruntime/modeling_seq2seq.py b/optimum/onnxruntime/modeling_seq2seq.py
@@ -1142,6 +1142,9 @@ class ORTModelForSeq2SeqLM(ORTModelForConditionalGeneration, GenerationMixin):
     auto_model_class = AutoModelForSeq2SeqLM
     main_input_name = "input_ids"
 
+    _supports_cache_class = False
+    _supports_static_cache = False
+
     def __init__(
         self,
         encoder_session: ort.InferenceSession,