Use Welford LayerNrom

intel · Feb 19, 2025 · 4358e06 · 4358e06
1 parent 6dcf059
commit 4358e06
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/ATen/native/xpu/sycl/LayerNormKernels.cpp b/src/ATen/native/xpu/sycl/LayerNormKernels.cpp
@@ -382,7 +382,7 @@ WelfordDataLN compute_stats(
   }
   // intra-warp reduction
   auto sg = item_id.get_sub_group();
-  int sgSize = static_cast<int>(sg.get_local_range()[1]);
+  int sgSize = static_cast<int>(sg.get_local_range()[0]);
   for (int offset = (sgSize >> 1); offset > 0; offset >>= 1) {
     WelfordDataLN wdB{
         sycl::shift_group_left(sg, wd.mean, offset),