diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc index 5a7b2cf16a1f8..ef5d3c7c2eee2 100644 --- a/paddle/phi/infermeta/unary.cc +++ b/paddle/phi/infermeta/unary.cc @@ -1889,6 +1889,11 @@ void MatrixPowerInferMeta(const MetaTensor& x, int n, MetaTensor* out) { "The Input(X) should have at least 2 dimensions. But " "received a %d dimension tensor.", n_dim)); + for (int i = 0; i < n_dim; ++i) + PADDLE_ENFORCE_NE( + dims[i], + 0, + phi::errors::InvalidArgument("The size of Input(X) should not be 0.")); PADDLE_ENFORCE_EQ(dims[n_dim - 2], dims[n_dim - 1], phi::errors::InvalidArgument( diff --git a/python/paddle/fluid/tests/unittests/test_matrix_power_op.py b/python/paddle/fluid/tests/unittests/test_matrix_power_op.py index 29f82b0350d65..7f26a7170191f 100644 --- a/python/paddle/fluid/tests/unittests/test_matrix_power_op.py +++ b/python/paddle/fluid/tests/unittests/test_matrix_power_op.py @@ -312,6 +312,10 @@ def test_errors(self): input = fluid.data(name="input_3", shape=[4, 5], dtype="float32") self.assertRaises(ValueError, paddle.linalg.matrix_power, input, 2) + # The size of input should not be 0 + input = fluid.data(name="input_4", shape=[1, 1, 0, 0], dtype="float32") + self.assertRaises(ValueError, paddle.linalg.matrix_power, input, 2) + class TestMatrixPowerSingularAPI(unittest.TestCase): def setUp(self):