From d1f8d0d21409e59c7c48f710c371c8d2bb67ce55 Mon Sep 17 00:00:00 2001
From: ooooo <3164076421@qq.com>
Date: Fri, 7 Feb 2025 08:40:58 +0000
Subject: [PATCH] fix func args

---
 python/paddle/jit/dy2static/pir_partial_program.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
index d245a631d77c0e..cfdf82988fe199 100644
--- a/python/paddle/jit/dy2static/pir_partial_program.py
+++ b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -774,7 +774,12 @@ def _get_scope(self, program_id=None, use_scope_cache=False):
     def _create_program(self, is_infer_mode=False) -> RunnableProgram:
         if is_infer_mode:
 
-            def pass_fn(forward_program, backward_program, program_name_attr):
+            def pass_fn(
+                forward_program,
+                backward_program,
+                program_name_attr,
+                whole_program,
+            ):
                 # common pass
                 pm = paddle.base.libpaddle.pir.PassManager()
                 paddle.base.libpaddle.pir.infer_symbolic_shape_pass(