disable tpp for un-verified models

huggingface · Jul 15, 2024 · e237f6f · e237f6f
1 parent 9ef6766
commit e237f6f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optimum/intel/ipex/modeling_base.py b/optimum/intel/ipex/modeling_base.py
@@ -110,8 +110,8 @@ def ipex_jit_trace(model, task, use_cache):
             sample_inputs.pop("past_key_values")
 
     # Use Tensor Processing Primitives to accelerate linear, see https://arxiv.org/abs/2104.05755.
-    # Only ipex >= 2.3.0 supports tpp.
-    if is_ipex_version(">=", "2.3.0"):
+    # Only ipex >= 2.3.0 supports tpp. The tpp is only verified for llm in generation tasks.
+    if is_ipex_version(">=", "2.3.0") and task in _IPEX_EXPORTED_GENERATION_TASKS:
         _enable_tpp()
     model = ipex.optimize(model.eval(), dtype=model.dtype, inplace=True)
     # Disable repack while jit tracing to reduce the memory