fix flag

PaddlePaddle · Aug 29, 2023 · 354042f · 354042f
1 parent 4d7ad68
commit 354042f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/test/prim/new_ir_prim/test_decomp_op.py b/test/prim/new_ir_prim/test_decomp_op.py
@@ -17,6 +17,7 @@
 import paddle
 from paddle import ir
 from paddle.decomposition import decompose
+from paddle.framework import core
 
 paddle.enable_static()
 
@@ -44,7 +45,9 @@ def test_build_op(self):
         y = newir_program.block().ops[-2].results()
         orig_shape = y[0].shape
         paddle.framework.set_flags({"FLAGS_enable_new_ir_api": True})
+        core._set_prim_forward_enabled(True)
         y_new = decompose(newir_program, y)
+        core._set_prim_forward_enabled(False)
         new_shape = y_new[0].shape
         assert (
             orig_shape == new_shape

diff --git a/test/prim/new_ir_prim/test_prim_program.py b/test/prim/new_ir_prim/test_prim_program.py
@@ -22,7 +22,6 @@
 from paddle.framework import core
 
 paddle.enable_static()
-paddle.framework.set_flags({"FLAGS_enable_new_ir_api": True})
 
 
 class TestPrimMode(unittest.TestCase):