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

                  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();
  }
};


  }
};

template <typename T>
__global__ void FindMovingAverageAbsMaxKernel(const T* in_state,
                                              const T* in_accum,
                                              const T* cur_scale, const T rate,
                                              T* out_state, T* out_accum,
                                              T* out_scale) {
  T state = rate * (*in_state) + T(1.0f);
  T accum = rate * (*in_accum) + (*cur_scale);
  *out_state = state;
  *out_accum = accum;
  *out_scale = accum / state;
}

template struct FindRangeAbsMaxFunctor<platform::CUDADeviceContext, float>;

template <typename T>
