!15648 [GraphKernel] negative axis in Squeeze expander.

From: @chenlei_autodiff
Reviewed-by: @gaoxiong1,@anyrenwei
Signed-off-by: @anyrenwei
This commit is contained in:
mindspore-ci-bot 2021-04-26 14:18:21 +08:00 committed by Gitee
commit 168c64b60d
1 changed files with 2 additions and 4 deletions

View File

@ -34,10 +34,8 @@ class Squeeze(Expander):
if not axis:
out_shape = [d for d in shape if d != 1]
else:
out_shape = []
for idx, dim in enumerate(shape):
if idx not in axis:
out_shape.append(dim)
ndim = len(shape)
out_shape = [shape[i] for i in range(ndim) if not (i in axis or (i - ndim) in axis)]
if not out_shape:
out_shape = [1]
return out_shape