diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index ae3d6121061a64..9da1abfb129b2d 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -6659,7 +6659,7 @@ def i0(x, name=None):
             Tensor(shape=[5], dtype=float32, place=Place(cpu), stop_gradient=True,
             [0.99999994 , 1.26606596 , 2.27958512 , 4.88079262 , 11.30192089])
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.i0(x)
     else:
         check_variable_and_dtype(x, "x", ["float32", "float64"], "i0")
diff --git a/test/legacy_test/test_i0_op.py b/test/legacy_test/test_i0_op.py
index 4ff7514752e0be..84694fa71ff08f 100644
--- a/test/legacy_test/test_i0_op.py
+++ b/test/legacy_test/test_i0_op.py
@@ -20,6 +20,7 @@
 
 import paddle
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 np.random.seed(100)
 paddle.seed(100)
@@ -40,10 +41,12 @@ class TestI0API(unittest.TestCase):
 
     def setUp(self):
         self.x = np.array(self.DATA).astype(self.DTYPE)
+        self.out_ref = output_i0(self.x)
         self.place = [paddle.CPUPlace()]
         if core.is_compiled_with_cuda():
             self.place.append(paddle.CUDAPlace(0))
 
+    @test_with_pir_api
     def test_api_static(self):
         def run(place):
             paddle.enable_static()
@@ -58,8 +61,7 @@ def run(place):
                     feed={"x": self.x},
                     fetch_list=[out],
                 )
-                out_ref = output_i0(self.x)
-                np.testing.assert_allclose(res[0], out_ref, rtol=1e-5)
+                np.testing.assert_allclose(res[0], self.out_ref, rtol=1e-5)
             paddle.disable_static()
 
         for place in self.place:
@@ -130,13 +132,14 @@ def init_config(self):
         self.target = output_i0(self.inputs['x'])
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
         self.check_grad(
             ['x'],
             'out',
             user_defined_grads=[ref_i0_grad(self.case, 1 / self.case.size)],
+            check_pir=True,
         )