Dtype compliance: native_batch_norm

Reviewed By: SS-JIA

Differential Revision: D48371009

fbshipit-source-id: 9b1b410438b817dd8e3ec97c9eb8f2c60c14a2ee

Author: manuelcandales

kernels/portable/cpu/op_native_batch_norm.cpp

@@ -44,7 +44,16 @@ std::tuple<Tensor&, Tensor&, Tensor&> _native_batch_norm_legit_no_training_out(
   ET_KERNEL_CHECK(
       ctx,
       check_batch_norm_args(
-          in, weight, bias, running_mean, running_var, momentum, eps, out),
+          in,
+          weight,
+          bias,
+          running_mean,
+          running_var,
+          momentum,
+          eps,
+          out,
+          mean_out,
+          var_out),
       InvalidArgument,
       ret_val);
   // For now, only support the default dim order

kernels/portable/cpu/util/normalization_ops_util.cpp

@@ -23,11 +23,15 @@ bool check_batch_norm_args(
     const Tensor& running_var,
     double momentum,
     double eps,
-    Tensor& out) {
+    Tensor& out,
+    Tensor& mean_out,
+    Tensor& var_out) {
   // All tensors must be the same dtype
   ET_LOG_AND_RETURN_IF_FALSE(
       tensors_have_same_dtype(in, running_mean, running_var));
   ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
+  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, mean_out));
+  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, var_out));
   if (weight.has_value()) {
     ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, weight.value()));
   }

kernels/portable/cpu/util/normalization_ops_util.h

@@ -21,7 +21,9 @@ bool check_batch_norm_args(
     const Tensor& running_var,
     double momentum,
     double eps,
-    Tensor& out);
+    Tensor& out,
+    Tensor& mean_out,
+    Tensor& var_out);
 
 bool check_layer_norm_args(
     const Tensor& input,