       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(!no_gamma){
+        if(need_gamma_grad) {
           local_gamma_grad[i] += (ele_x - local_mean[i]) / local_std[i] * ele_og;
         }
-        if(!no_beta) {
+        if(need_beta_grad) {
           local_beta_grad[i] += ele_og;
         }
       }
     }
     // Write to destination
     for(int i = 0; i < LOAD_UNROLL; ++i) {
-      if(!no_gamma) {
-        gamma_grad[l + i] += local_gamma_grad[i];
-      }
-      if(!no_beta) {
-        beta_grad[l + i] += local_beta_grad[i];
+      if(need_gamma_grad) {
+        if(gamma_addto) {
+          gamma_grad[l + i] += local_gamma_grad[i];
+        }
+      }
+      if(need_beta_grad) {
+        if(beta_addto) {
+          beta_grad[l + i] += local_beta_grad[i];
+        }
       }
     }
   }

  CHECK_EQ(mean_data.CheckContiguous(), true);
  CHECK_EQ(std_data.CheckContiguous(), true);

  // Lauch the kernel. The dynamic shared memory size is sizeof(DType) * threadDim.y + sizeof(DType) *
  int nbatch = data_shape[0];
  int nchannel = data_shape[1];
  float eps = param.eps;
  int ngrid_x = (nbatch > kMaxGridDim) ? (nbatch + kBaseGridNum - 1) / kBaseGridNum : nbatch;
  int ngrid_y = (nbatch > kMaxGridDim) ? kBaseGridNum : 1;
  int nthread_y = 0;
  const dim3 dimGrid(ngrid_x, ngrid_y, 1);
  if(nchannel <= 32) {
    nthread_y = 1;
  } else {
    nthread_y = 2;
  }
  cudaStream_t stream = Stream<gpu>::GetStream(ctx.get_stream<gpu>());
  const dim3 dimBlock(32, nthread_y, 1);
  MSHADOW_REAL_TYPE_SWITCH(in_data.type_flag_, DType, {
    int nshared = nthread_y > 1 ? nthread_y * sizeof(DType) + (nthread_y / 2) * sizeof(DType) : 0;
    CheckLaunchParam(dimGrid, dimBlock);

    // Try to use the accelerated CUDA kernels
    return LayerNormGradGPUContig(param, ctx, inputs, req, outputs);
  }
  return LayerNormGradComputeGeneral<gpu>(attrs, ctx, inputs, req, outputs);
}

