     }
     // Summation
     for(int b = 0; b < nbatch; ++b) {
+      DType local_mean = 0;
+      DType local_std = 0;
+      if(need_gamma_grad) {
+        local_mean = mean_data[b];
+        local_std = std_data[b];
+      }
 #pragma unroll
       for(int i = 0; i < LOAD_UNROLL; ++i) {
         DType ele_og = out_grad[b * nchannel + l + i];
         DType ele_x = in_data[b * nchannel + l + i];
         if(need_gamma_grad) {
-          local_gamma_grad[i] += (ele_x - local_mean[i]) / local_std[i] * ele_og;
+          local_gamma_grad[i] += (ele_x - local_mean) / local_std * ele_og;
         }
         if(need_beta_grad) {
           local_beta_grad[i] += ele_og;

  }
  // Handle the rest elements that cannot be divided by LOAD_UNROLL
  for(; l < nchannel; ++l) {
    local_mean[0] = mean_data[l];
    local_std[0] = std_data[l];
    local_gamma_grad[0] = 0;
    local_beta_grad[0] = 0;
    for(int b = 0; b < nbatch; ++b) {
      DType ele_og = out_grad[b * nchannel + l];
      DType ele_x = in_data[b * nchannel + l];
      if(need_gamma_grad) {
        local_gamma_grad[0] += (ele_x - local_mean[0]) / local_std[0] * ele_og;
      }
      if(need_beta_grad) {
        local_beta_grad[0] += ele_og;

    // Initialize the shared memory
#pragma unroll
    for(int i = 0; i < LOAD_UNROLL; i++) {
      if(need_gamma_grad) {
        local_gamma_grad[i] = 0;
      }
