[ET][Portable] Fix op native_batch_norm

Resize `mean_out` & `var_out`

Differential Revision: [D50081049](https://our.internmc.facebook.com/intern/diff/D50081049/)

ghstack-source-id: 203404908
Pull Request resolved: #725

Author: manuelcandales

kernels/portable/cpu/op_native_batch_norm.cpp

@@ -41,6 +41,12 @@ std::tuple<Tensor&, Tensor&, Tensor&> _native_batch_norm_legit_no_training_out(
       InvalidArgument,
       ret_val);
 
+  ET_KERNEL_CHECK(
+      ctx, resize_tensor(mean_out, {0}) == Error::Ok, InvalidArgument, ret_val);
+
+  ET_KERNEL_CHECK(
+      ctx, resize_tensor(var_out, {0}) == Error::Ok, InvalidArgument, ret_val);
+
   ET_KERNEL_CHECK(
       ctx,
       check_batch_norm_args(
@@ -56,6 +62,7 @@ std::tuple<Tensor&, Tensor&, Tensor&> _native_batch_norm_legit_no_training_out(
           var_out),
       InvalidArgument,
       ret_val);
+
   // For now, only support the default dim order
   ET_KERNEL_CHECK(
       ctx,

kernels/portable/cpu/util/normalization_ops_util.cpp

@@ -45,6 +45,9 @@ bool check_batch_norm_args(
   ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(running_mean, 1));
   ET_LOG_AND_RETURN_IF_FALSE(
       tensors_have_same_size_at_dims(running_mean, 0, in, C_dim));
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(running_var, 1));
+  ET_LOG_AND_RETURN_IF_FALSE(
+      tensors_have_same_size_at_dims(running_var, 0, in, C_dim));
   if (weight.has_value()) {
     ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(weight.value(), 1));
     ET_LOG_AND_RETURN_IF_FALSE(