未验证 提交 f402d8d8 编写于 作者: L Leo Chen 提交者: GitHub

fix bug when axis is a tensor with more than 1 element (#27263)

上级 20a84820
...@@ -6309,7 +6309,7 @@ def unsqueeze(input, axes, name=None): ...@@ -6309,7 +6309,7 @@ def unsqueeze(input, axes, name=None):
if isinstance(axes, int): if isinstance(axes, int):
axes = [axes] axes = [axes]
elif isinstance(axes, Variable): elif isinstance(axes, Variable):
axes = [axes.numpy().item(0)] axes = axes.numpy().tolist()
elif isinstance(axes, (list, tuple)): elif isinstance(axes, (list, tuple)):
axes = [ axes = [
item.numpy().item(0) if isinstance(item, Variable) else item item.numpy().item(0) if isinstance(item, Variable) else item
......
...@@ -167,8 +167,9 @@ class API_TestDyUnsqueezeAxisTensor(unittest.TestCase): ...@@ -167,8 +167,9 @@ class API_TestDyUnsqueezeAxisTensor(unittest.TestCase):
with fluid.dygraph.guard(): with fluid.dygraph.guard():
input1 = np.random.random([5, 10]).astype("int32") input1 = np.random.random([5, 10]).astype("int32")
out1 = np.expand_dims(input1, axis=1) out1 = np.expand_dims(input1, axis=1)
out1 = np.expand_dims(out1, axis=2)
input = fluid.dygraph.to_variable(input1) input = fluid.dygraph.to_variable(input1)
output = paddle.unsqueeze(input, axis=paddle.to_tensor([1])) output = paddle.unsqueeze(input, axis=paddle.to_tensor([1, 2]))
out_np = output.numpy() out_np = output.numpy()
self.assertTrue(np.array_equal(out1, out_np)) self.assertTrue(np.array_equal(out1, out_np))
self.assertEqual(out1.shape, out_np.shape) self.assertEqual(out1.shape, out_np.shape)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册