fix bug for fp32 batchnorm_op when using nhwc data_layout (PaddlePadd…

…le#37020)
GuoxiaWang · Dec 27, 2021 · bd393a9 · bd393a9
1 parent 462ee10
commit bd393a9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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);