From f402d8d8c0239f39f67342e09ba865d4d419a0ce Mon Sep 17 00:00:00 2001 From: Leo Chen Date: Fri, 11 Sep 2020 18:22:57 +0800 Subject: [PATCH] fix bug when axis is a tensor with more than 1 element (#27263) --- python/paddle/fluid/layers/nn.py | 2 +- python/paddle/fluid/tests/unittests/test_unsqueeze_op.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py index 5a14b9fdc7b..bc9f182d95e 100755 --- a/python/paddle/fluid/layers/nn.py +++ b/python/paddle/fluid/layers/nn.py @@ -6309,7 +6309,7 @@ def unsqueeze(input, axes, name=None): if isinstance(axes, int): axes = [axes] elif isinstance(axes, Variable): - axes = [axes.numpy().item(0)] + axes = axes.numpy().tolist() elif isinstance(axes, (list, tuple)): axes = [ item.numpy().item(0) if isinstance(item, Variable) else item diff --git a/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py b/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py index 6f713172f1b..1975e430602 100644 --- a/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py +++ b/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py @@ -167,8 +167,9 @@ class API_TestDyUnsqueezeAxisTensor(unittest.TestCase): with fluid.dygraph.guard(): input1 = np.random.random([5, 10]).astype("int32") out1 = np.expand_dims(input1, axis=1) + out1 = np.expand_dims(out1, axis=2) 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() self.assertTrue(np.array_equal(out1, out_np)) self.assertEqual(out1.shape, out_np.shape) -- GitLab