add cinn flag

cyber-pioneer · Mar 6, 2024 · 93ced7e · 93ced7e
1 parent 5d90acd
commit 93ced7e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/llm/predictor.py b/llm/predictor.py
@@ -26,7 +26,7 @@
 import paddle
 import paddle.distributed.fleet.base.topology as tp
 import paddle.incubate.multiprocessing as mp
-from paddle.base.framework import in_pir_executor_mode
+from paddle.base.framework import in_cinn_mode, in_pir_executor_mode
 from paddle.distributed import fleet
 from utils import (
     dybatch_preprocess,
@@ -363,6 +363,8 @@ def __init__(self, config: PredictorArgument, tokenizer: PretrainedTokenizer = N
         inference_config.enable_new_executor()
         if in_pir_executor_mode():
             inference_config.enable_new_ir()
+            if in_cinn_mode():
+                inference_config.enable_cinn()
 
         with static_mode_guard():
             self.predictor = paddle.inference.create_predictor(inference_config)