Fix F.unpool1d when output_size is tuple 易用性提升 (#65910)

NKNaN · web-flow · commit bcc26529a9b5 · 2024-07-18T17:11:40.000+08:00
* fix bug

* fix bug
diff --git a/python/paddle/nn/functional/pooling.py b/python/paddle/nn/functional/pooling.py
@@ -833,7 +833,11 @@ def max_unpool1d(
     padding = _expand_low_nd_padding(padding)
 
     if output_size is not None:
-        output_size = output_size[:2] + [1] + output_size[2:]
+        output_size = (
+            output_size[:2]
+            + ([1] if isinstance(output_size, list) else (1,))
+            + output_size[2:]
+        )
     output_size = _unpool_output_size(
         x, kernel_size, stride, padding, output_size
     )
diff --git a/test/legacy_test/test_unpool1d_op.py b/test/legacy_test/test_unpool1d_op.py
@@ -164,6 +164,35 @@ def test_case(self):
         paddle.enable_static()
 
 
+class TestUnpool1DOpAPI_dygraph5(unittest.TestCase):
+    def test_case(self):
+        places = [paddle.CPUPlace()]
+        if paddle.base.core.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        for place in places:
+            paddle.disable_static()
+            input_data = np.arange(3 * 16).reshape([1, 3, 16]).astype("float32")
+            input_x = paddle.to_tensor(input_data)
+            output, indices = F.max_pool1d(
+                input_x, kernel_size=2, stride=2, return_mask=True
+            )
+            output_unpool = F.max_unpool1d(
+                output.astype("int64"),
+                indices,
+                kernel_size=2,
+                stride=2,
+                output_size=tuple(input_x.shape),
+            )
+            expected_output_unpool = unpool1dmax_forward_naive(
+                output.numpy(), indices.numpy(), [2], [2], [0], [16]
+            )
+            np.testing.assert_allclose(
+                output_unpool.numpy(), expected_output_unpool, rtol=1e-05
+            )
+
+        paddle.enable_static()
+
+
 class TestUnpool1DOpAPI_static(unittest.TestCase):
     @test_with_pir_api
     def test_case(self):
