fix rank=1 (#51413)

PaddlePaddle · Mar 14, 2023 · b4f49aa · b4f49aa
1 parent b642461
commit b4f49aa
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/paddle/phi/infermeta/ternary.cc b/paddle/phi/infermeta/ternary.cc
@@ -1073,7 +1073,7 @@ void ScatterNdAddInferMeta(const MetaTensor& x,
             index_dims[index_dims_size - 1],
             ref_dims_size));
     PADDLE_ENFORCE_GE(index_dims_size,
-                      2UL,
+                      1UL,
                       phi::errors::InvalidArgument(
                           "The rank of Input(Index) should be greater than 1, "
                           "but received the rank of Input(Index) is %d.",

diff --git a/python/paddle/fluid/tests/unittests/test_gather_nd_op.py b/python/paddle/fluid/tests/unittests/test_gather_nd_op.py
@@ -54,7 +54,7 @@ def test_check_output(self):
         self.check_output(check_eager=False)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_eager=False)
+        self.check_grad(['X'], 'Out', check_eager=False, check_prim=True)
 
 
 class TestGatherNdOpWithLowIndex(OpTest):