   int tid = threadIdx.x;
   int bid = blockIdx.x;
   const T* in_current = in + tid * cout_wh_size + bid * wh_size;
-  shared_max_data[tid] = T(0);
+  T local_max_data = T(0);
   for (int i = 0; i < wh_size; i++) {
     T tmp = fabs(in_current[i]);
-    if (tmp > shared_max_data[tid]) {
-      shared_max_data[tid] = tmp;
-    }
-  }
+    if (tmp > local_max_data) {
+      local_max_data = tmp;
+    }
+  }
+  shared_max_data[tid] = local_max_data;
   __syncthreads();
 
   int len = blockDim.x;

                  framework::Tensor* out_accum, framework::Tensor* out_scale) {
    const auto gpu_place = ctx.GetPlace();

    T accum;
    T state;
    T scale;
    memory::Copy(platform::CPUPlace(), &accum, gpu_place, in_accum.data<T>(),
                 sizeof(T), ctx.stream());
    memory::Copy(platform::CPUPlace(), &state, gpu_place, in_state.data<T>(),
                 sizeof(T), ctx.stream());
    memory::Copy(platform::CPUPlace(), &scale, gpu_place, cur_scale, sizeof(T),
                 ctx.stream());
    ctx.Wait();

    T rate_t = static_cast<T>(rate);
    state = rate_t * state + static_cast<T>(1.0);
    accum = rate_t * accum + scale;
    scale = accum / state;

    memory::Copy(gpu_place, out_accum->mutable_data<T>(gpu_place),
                 platform::CPUPlace(), &accum, sizeof(T), ctx.stream());
    memory::Copy(gpu_place, out_state->mutable_data<T>(gpu_place),
                 platform::CPUPlace(), &state, sizeof(T), ctx.stream());
    memory::Copy(gpu_place, out_scale->mutable_data<T>(gpu_place),
                 platform::CPUPlace(), &scale, sizeof(T), ctx.stream());
    ctx.Wait();
  }
};


  int channel_size = n / c;
  const T* in_c = in + blockIdx.x * channel_size;
  extern __shared__ T shared_max_data[];
  T local_max_data = T(0);
  for (int i = tid; i < channel_size; i += blockDim.x) {
    T tmp = fabs(in_c[i]);
    if (tmp > local_max_data) {
      local_max_data = tmp;
    }
  }
  shared_max_data[tid] = local_max_data;
  __syncthreads();
  for (int i = blockDim.x / 2; i > 0; i >>= 1) {
    if (tid < i && (shared_max_data[tid] < shared_max_data[tid + i])) {
