【CPU】change Batch norm reserve_space dtype (#72314)

xiaoguoguo626807 · web-flow · commit 1f6403ddc48f · 2025-04-17T14:05:03.000+08:00
* batch_norm nouse output

* Apply suggestions from code review
diff --git a/paddle/phi/kernels/cpu/batch_norm_kernel.cc b/paddle/phi/kernels/cpu/batch_norm_kernel.cc
@@ -101,8 +101,8 @@ void BatchNormKernel(const Context& ctx,
         ctx.template Alloc<T>(saved_variance), C);
     saved_mean_e.setZero();
     saved_variance_e.setZero();
-    EigenVectorArrayMap<T> reserve_space_e(ctx.template Alloc<T>(reserve_space),
-                                           0);
+    EigenVectorArrayMap<uint8_t> reserve_space_e(
+        ctx.template Alloc<uint8_t>(reserve_space), 0);
     reserve_space_e.setZero();
 
     EigenVectorArrayMap<T> running_mean_arr(ctx.template Alloc<T>(mean_out), C);
@@ -222,4 +222,6 @@ void BatchNormKernel(const Context& ctx,
 }  // namespace phi
 
 PD_REGISTER_KERNEL(
-    batch_norm, CPU, ALL_LAYOUT, phi::BatchNormKernel, float, double) {}
+    batch_norm, CPU, ALL_LAYOUT, phi::BatchNormKernel, float, double) {
+  kernel->OutputAt(5).SetDataType(phi::DataType::UINT8);
+}
diff --git a/paddle/phi/kernels/xpu/batch_norm_kernel.cc b/paddle/phi/kernels/xpu/batch_norm_kernel.cc
@@ -98,7 +98,6 @@ void BatchNormKernel(const Context& dev_ctx,
   dev_ctx.template Alloc<float>(variance_out);
   dev_ctx.template Alloc<float>(saved_mean);
   dev_ctx.template Alloc<float>(saved_variance);
-
   PADDLE_ENFORCE_LE(
       x_dims.size(),
       5,
@@ -164,4 +163,5 @@ PD_REGISTER_KERNEL(batch_norm,
   kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
   kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
   kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(5).SetDataType(phi::DataType::UINT8);
 }
