From 9dde26f656f886307818ee9dd479030cf8d1ba05 Mon Sep 17 00:00:00 2001
From: xiaoting <31891223+tink2123@users.noreply.github.com>
Date: Tue, 27 Dec 2022 16:45:13 +0800
Subject: [PATCH] fix fold for large bs (#49337)

* fix fold for large bs

* fix fold for large bs
---
 paddle/phi/kernels/impl/fold_grad_kernel_impl.h    |  7 ++-----
 paddle/phi/kernels/impl/fold_kernel_impl.h         |  8 +++-----
 .../paddle/fluid/tests/unittests/test_fold_op.py   | 14 ++++++++++++++
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/paddle/phi/kernels/impl/fold_grad_kernel_impl.h b/paddle/phi/kernels/impl/fold_grad_kernel_impl.h
index 7118de3174f..3143106edc2 100644
--- a/paddle/phi/kernels/impl/fold_grad_kernel_impl.h
+++ b/paddle/phi/kernels/impl/fold_grad_kernel_impl.h
@@ -54,11 +54,8 @@ void FoldGradKernel(const Context& ctx,
 
   DDim out_shape =
       make_ddim({n_output_plane, output_sizes[0], output_sizes[1]});
-  DDim input_matrix_shape = make_ddim({x_dims[0],
-                                       kernel_sizes[0],
-                                       kernel_sizes[1],
-                                       output_height,
-                                       output_width});
+  DDim input_matrix_shape = make_ddim(
+      {1, kernel_sizes[0], kernel_sizes[1], output_height, output_width});
 
   phi::funcs::Im2ColFunctor<phi::funcs::ColFormat::kCFO, Context, T> im2col;
 
diff --git a/paddle/phi/kernels/impl/fold_kernel_impl.h b/paddle/phi/kernels/impl/fold_kernel_impl.h
index 21864b00cae..694d754ecfb 100644
--- a/paddle/phi/kernels/impl/fold_kernel_impl.h
+++ b/paddle/phi/kernels/impl/fold_kernel_impl.h
@@ -54,11 +54,8 @@ void FoldKernel(const Context& ctx,
   DDim output_shape =
       make_ddim({n_output_plane, output_sizes[0], output_sizes[1]});
 
-  DDim input_matrix_shape = make_ddim({x_dims[0],
-                                       kernel_sizes[0],
-                                       kernel_sizes[1],
-                                       output_height,
-                                       output_width});
+  DDim input_matrix_shape = make_ddim(
+      {1, kernel_sizes[0], kernel_sizes[1], output_height, output_width});
 
   phi::funcs::SetConstant<Context, T> set_zero;
   set_zero(ctx, out, static_cast<T>(0));
@@ -66,6 +63,7 @@ void FoldKernel(const Context& ctx,
   for (int i = 0; i < batch_size; i++) {
     DenseTensor out_batch =
         out->Slice(i, i + 1).Resize(output_shape);  // im size=3
+
     DenseTensor in_batch =
         x.Slice(i, i + 1).Resize(input_matrix_shape);  // col size=5
     col2im(ctx, in_batch, dilations, strides, paddings, &out_batch);
diff --git a/python/paddle/fluid/tests/unittests/test_fold_op.py b/python/paddle/fluid/tests/unittests/test_fold_op.py
index 71f69aa6d67..1f3193fa1fd 100644
--- a/python/paddle/fluid/tests/unittests/test_fold_op.py
+++ b/python/paddle/fluid/tests/unittests/test_fold_op.py
@@ -130,6 +130,20 @@ class TestFoldOp(OpTest):
         self.check_grad(['X'], 'Y', check_eager=True)
 
 
+class TestFoldshape(TestFoldOp):
+    def init_data(self):
+        self.batch_size = 8
+        self.input_channels = 3 * 3 * 3
+        self.length = 6
+        self.kernel_sizes = [3, 3]
+        self.strides = [1, 1]
+        self.paddings = [0, 0, 0, 0]
+        self.dilations = [1, 1]
+        self.output_sizes = [4, 5]
+        input_shape = [self.batch_size, self.input_channels, self.length]
+        self.x = np.random.rand(*input_shape).astype(np.float64)
+
+
 class TestFoldAPI(TestFoldOp):
 
     # This is for test on paddle.nn.Fold
-- 
GitLab