Added check if reuse_cache supported by selected model_type (huggingf…

…ace#675)
huijuanzh · Feb 2, 2024 · 4077355 · 4077355
1 parent 93264ba
commit 4077355
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py
@@ -600,6 +600,7 @@ def generate(
         if model_kwargs["reduce_recompile"]:
             assert generation_config.bucket_size
         if generation_config.reuse_cache:
+            assert self.config.model_type in ["llama"], "reuse_cache only supported by llama at the moment"
             assert generation_config.bucket_size <= 0, "reuse_cache and bucketing flags set together"
 
         if generation_config.static_shapes: