     //    d_seq[i] += d_cols[i+1, 1]
     //    d_seq[i] += d_cols[i+2, 0]
 
-    for (int b = blockIdx.x * blockDim.x + threadIdx.x; b < B; b += blockDim.x * gridDim.x)
-    {
-        int nF = nW * 2 + 1;
-        int seq_row = b * I;
+    int _loop_start = blockIdx.x * blockDim.x + threadIdx.x;
+    int _loop_stride = blockDim.x * gridDim.x;
+    int nF = nW * 2 + 1;
+    int end_d_cols = B * I * nF;
+    for (int b = _loop_start; b < B; b += _loop_stride)
+    {
+        float* d_seqs_b = &d_seqs[b*I];
         int col_feat = nF * I;
         for (int f=-nW; f < (nW+1); ++f)
         {
-            int col_row = (b+f) * (I * nF);
+            int col_row = (b+f) * (I*nF);
             col_feat -= I;
-            if ((col_row >= 0) && (col_row < (B*I*nF)))
+            if ((col_row >= 0) && (col_row < end_d_cols))
             {
                 int start = col_row + col_feat;
-                if ((start >= 0) && ((start+I) < (B*I*nF)))
+                if ((start >= 0) && ((start+I) < end_d_cols))
                 {
                     for (int i=0; i < I; ++i)
-            	    atomicAdd(&d_seqs[seq_row+i], d_cols[start+i]);
+                        d_seqs_b[i] += d_cols[start+i];
                 }
             }
         }

void backprop_max_pool(float* dX,
    const float* d_maxes, const int* which, const int* lengths, int B, int T, int O)
{
    for (int row = blockIdx.x * blockDim.x + threadIdx.x; row < T; row += blockDim.x * gridDim.x)
    {
        // Find the sequence item we're working on
        int seq_start = 0;
        int b = 0;
        while ((b < B) && (seq_start+lengths[b]) < row)
        {
           seq_start += lengths[b];
           b += 1;
        }

        dX = &dX[row];
        which = &which[b];
        d_maxes = &d_maxes[b];
 
        for (int j=0; j < O; ++j)
        {
            // If we used the value for this cell,
            // pass the gradient
            if (which[j] == (row-seq_start))
               dX[j] = d_maxes[j];
        }
    }
}



extern "C" __global__
void max_pool(float* maxes, int* which,
    const float* X, const int* lengths, int B, int T, int O)
{
    int _loop_start = blockIdx.x * blockDim.x + threadIdx.x;
    int _loop_stride = blockDim.x * gridDim.x;
    for (int b = _loop_start; b < B; b += _loop_stride)
    {

        // Go to the regions we're working on
        float* maxes_b = &maxes[b*O];
        float* which_b = &which[b*O];
        // Find the sequence item we're working on
        int seq_start = 0;
        for (int i=0; i < b; ++i) {
            seq_start += lengths[b];
        }
        // Each invocation of the kernel maxes one sequence.
        // Start by assuming maxes are at i=0
        for (int i=0; i < O; ++i) {
            maxes_b[i] = X[i];
            which_b[i] = 0;
        }
        
        int length = lengths[b];
        for (int i=1; i < length; ++i) // Iterate over rows
        {
            float* X_t = &X[(seq_start+i)*O]
            for (int j=0; j < O; ++j)
            {
                if (X_t[j] > maxes_b[j])
                {
                    maxes_b[j] = X_t[j];
                    which_b[j] = i;
                }
            }
        }

    }
}

