diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc index 5a7b2cf16a1f8c..7827fd1ee87e95 100644 --- a/paddle/phi/infermeta/unary.cc +++ b/paddle/phi/infermeta/unary.cc @@ -3600,6 +3600,10 @@ void SplitWithNumInferMeta(const MetaTensor& x, auto input_axis_dim = x.dims().at(axis_value); // step1: get formated sections std::vector sections_vec; + PADDLE_ENFORCE_NE( + num, + 0, + phi::errors::InvalidArgument("Attr(num_or_sections) should not be 0.")); PADDLE_ENFORCE_EQ(input_axis_dim % num, 0, phi::errors::InvalidArgument( diff --git a/python/paddle/fluid/tests/unittests/test_split_op.py b/python/paddle/fluid/tests/unittests/test_split_op.py index 943861ab2f57b3..962966a9c8fc73 100644 --- a/python/paddle/fluid/tests/unittests/test_split_op.py +++ b/python/paddle/fluid/tests/unittests/test_split_op.py @@ -336,6 +336,14 @@ def test_axis_type_tensor(): self.assertRaises(TypeError, test_axis_type_tensor) + with paddle.fluid.dygraph.guard(): + + def test_0_num_tensor(): + x = paddle.uniform([1, 1, 1], dtype='float32') + paddle.split(x, num_or_sections=0) + + self.assertRaises(ValueError, test_0_num_tensor) + class API_TestSplit(unittest.TestCase): def test_out(self):