forked from mindspore-Ecosystem/mindspore
!15648 [GraphKernel] negative axis in Squeeze expander.
From: @chenlei_autodiff Reviewed-by: @gaoxiong1,@anyrenwei Signed-off-by: @anyrenwei
This commit is contained in:
commit
168c64b60d
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue