From 8ef7102ce7fbcc95c062d1188429cbb89224aa33 Mon Sep 17 00:00:00 2001
From: Guoxia Wang <mingzilaochongtu@gmail.com>
Date: Wed, 29 Dec 2021 16:23:47 +0800
Subject: [PATCH] fix bug for fp32 batchnorm_op when using nhwc data_layout
 (#37020) (#38497)

---
 paddle/fluid/operators/batch_norm_op.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/paddle/fluid/operators/batch_norm_op.cu b/paddle/fluid/operators/batch_norm_op.cu
index 3d26c2c5708..b4cf9c48df2 100644
--- a/paddle/fluid/operators/batch_norm_op.cu
+++ b/paddle/fluid/operators/batch_norm_op.cu
@@ -916,7 +916,7 @@ class BatchNormGradKernel<platform::CUDADeviceContext, T>
     Tensor transformed_d_y(d_y->type());
     Tensor transformed_d_x;
     if (data_layout == DataLayout::kNHWC &&
-        compute_format == DataLayout::kNCHW) {
+        compute_format == DataLayout::kNCHW && x_dims.size() > 2) {
       VLOG(3) << "Transform input tensor from NHWC to NCHW.";
       ResizeToChannelFirst<platform::CUDADeviceContext, T>(ctx, x,
                                                            &transformed_x);
-- 
GitLab