       float* so = out + j * cols;
       const float* sp = in + j * cols;
 
-      float* _sum = _share + blockDim.x;
+      float* _sum = _share;
       _sum[threadIdx.x] = 0.0f;
       for(int tid = 0; tid < cols; tid += blockDim.x) {
         int id = tid + threadIdx.x;

    int j = bid + blockIdx.x;
    if(j < rows) {
      extern __shared__ float _share[];
      float* _sum = _share + blockDim.x;

      float* gradRow = grad + j * cols;
      const float* adjRow = adj + j * cols;


  int blocks = std::min(MAX_BLOCKS, (int)m);
  int threads = std::min(MAX_THREADS, (int)k);
  int shared = sizeof(float) * threads;

  gAtt<<<blocks, threads, shared>>>(
      out->data(), va->data(), context->data(), state->data(), m, k, b, t);
