Update

[ghstack-poisoned]

Author: swolchok

kernels/portable/cpu/op_mean.cpp

@@ -45,22 +45,23 @@ Tensor& mean_dim_out(
       out);
 
   ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, "mean.out", CTYPE_IN, [&] {
-    ET_SWITCH_FLOATHBF16_TYPES(out.scalar_type(), ctx, "mean.out", CTYPE_OUT, [&] {
-      CTYPE_OUT* out_data = out.mutable_data_ptr<CTYPE_OUT>();
-      const size_t num = get_reduced_dim_product(in, dim_list);
-      for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {
-        CTYPE_OUT sum = 0;
-        if (in.numel() > 0) {
-          sum = map_reduce_over_dim_list<CTYPE_IN, CTYPE_OUT>(
-              [](CTYPE_IN v) { return static_cast<CTYPE_OUT>(v); },
-              [](CTYPE_OUT outv, CTYPE_OUT acc) { return acc + outv; },
-              in,
-              dim_list,
-              out_ix);
-        }
-        out_data[out_ix] = sum / static_cast<float>(num);
-      }
-    });
+    ET_SWITCH_FLOATHBF16_TYPES(
+        out.scalar_type(), ctx, "mean.out", CTYPE_OUT, [&] {
+          CTYPE_OUT* out_data = out.mutable_data_ptr<CTYPE_OUT>();
+          const size_t num = get_reduced_dim_product(in, dim_list);
+          for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {
+            CTYPE_OUT sum = 0;
+            if (in.numel() > 0) {
+              sum = map_reduce_over_dim_list<CTYPE_IN, CTYPE_OUT>(
+                  [](CTYPE_IN v) { return static_cast<CTYPE_OUT>(v); },
+                  [](CTYPE_OUT outv, CTYPE_OUT acc) { return acc + outv; },
+                  in,
+                  dim_list,
+                  out_ix);
+            }
+            out_data[out_ix] = sum / static_cast<float>(num);
+          }
+        });
   });
 
   return out;