fix bug of batch_norm_grad kernel with fp16

Author: XiaoguangHu01

paddle/phi/kernels/gpu/batch_norm_grad_kernel.cu

@@ -988,10 +988,9 @@ PD_REGISTER_KERNEL(batch_norm_grad,
                    double,
                    phi::dtype::float16) {
   if (kernel_key.dtype() == phi::DataType::FLOAT16) {
-    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
-    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
-    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
-    kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(0).SetDataType(phi::DataType::FLOAT32);  //x_grad
+    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);  //scale_grad
+    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);  //bias_grad
   }
 }
 
@@ -1003,10 +1002,9 @@ PD_REGISTER_KERNEL(batch_norm_grad_raw,
                    double,
                    phi::dtype::float16) {
   if (kernel_key.dtype() == phi::DataType::FLOAT16) {
-    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
-    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
-    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
-    kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(0).SetDataType(phi::DataType::FLOAT32);  //x_grad
+    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);  //scale_grad
+    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);  //bias_grad
   }
 }