diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py index 32f7bf373bbbd8dbe59befad263fae89ecaeeca3..95f8fa449bd5ffb3eee71449930d0701a8afc1e7 100644 --- a/python/paddle/tensor/search.py +++ b/python/paddle/tensor/search.py @@ -399,7 +399,7 @@ def nonzero(x, as_tuple=False): for i in range(rank): list_out.append( layers.slice( - outs, axes=[rank - 1], starts=[i], ends=[i + 1])) + outs, axes=[1], starts=[i], ends=[i + 1])) return tuple(list_out)