diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py index 042802a8e343bc..85595ba7668e6c 100755 --- a/python/paddle/tensor/search.py +++ b/python/paddle/tensor/search.py @@ -607,7 +607,7 @@ def mode(x, axis=-1, keepdim=False, name=None): [2, 1]])) """ - if in_dynamic_mode(): + if in_dynamic_or_pir_mode(): return _C_ops.mode(x, axis, keepdim) else: helper = LayerHelper("mode", **locals()) diff --git a/test/legacy_test/test_mode_op.py b/test/legacy_test/test_mode_op.py index 4f13158bb664dd..1ce3c270da4874 100644 --- a/test/legacy_test/test_mode_op.py +++ b/test/legacy_test/test_mode_op.py @@ -20,6 +20,7 @@ import paddle from paddle import base from paddle.base import core +from paddle.pir_utils import test_with_pir_api def _mode1D(a): @@ -112,12 +113,12 @@ def init_numeric_grads(self): def test_check_output(self): paddle.enable_static() - self.check_output() + self.check_output(check_pir=True) def test_check_grad(self): paddle.enable_static() grad = self.init_numeric_grads() - self.check_grad({'X'}, 'Out', user_defined_grads=[grad]) + self.check_grad({'X'}, 'Out', user_defined_grads=[grad], check_pir=True) @unittest.skipIf( @@ -148,7 +149,7 @@ def test_check_output(self): place = core.CUDAPlace(0) paddle.enable_static() if core.is_bfloat16_supported(place): - self.check_output_with_place(place) + self.check_output_with_place(place, check_pir=True) def test_check_grad(self): place = core.CUDAPlace(0) @@ -157,7 +158,7 @@ def test_check_grad(self): if core.is_bfloat16_supported(place): self.check_grad_with_place( - place, {'X'}, 'Out', user_defined_grads=[grad] + place, {'X'}, 'Out', user_defined_grads=[grad], check_pir=True ) @@ -243,6 +244,7 @@ def setUp(self): np.random.random((2, 10, 10)) * 1000, dtype=np.float64 ) + @test_with_pir_api def test_run_static(self): paddle.enable_static() with paddle.static.program_guard(