From ea02bf6e7af6291808e841823451d8095cea5482 Mon Sep 17 00:00:00 2001 From: tom__chen Date: Mon, 29 Mar 2021 11:28:17 -0400 Subject: [PATCH] index_add GPU op check rank of indices tensor --- mindspore/ops/operations/math_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mindspore/ops/operations/math_ops.py b/mindspore/ops/operations/math_ops.py index 6ec8c5ea8f..0735648d56 100644 --- a/mindspore/ops/operations/math_ops.py +++ b/mindspore/ops/operations/math_ops.py @@ -4601,6 +4601,7 @@ class IndexAdd(PrimitiveWithInfer): validator.check("x rank", len(x_shape), "y rank", len(y_shape), Rel.EQ, self.name) x_rank = len(x_shape) validator.check_int_range(self.axis, -x_rank - 1, x_rank, Rel.INC_NEITHER, 'axis', self.name) + validator.check_equal_int(len(idx_shape), 1, "rank of idx_shape", self.name) validator.check("size of indices", idx_shape[0], "dimension of y[axis]", y_shape[self.axis], Rel.EQ, self.name) axis = self.axis if self.axis >= 0 else x_rank + self.axis