 void backprop_sum_pool(float* dX, const float* d_sum, 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)
+    int _loop_start = blockIdx.x * blockDim.x + threadIdx.x;
+    int _loop_stride = blockDim.x * gridDim.x;
+    int seq_start = 0;
+    int b = 0;
+    for (int t = _loop_start; t < T; t += _loop_stride)
     { 
         // Find the sequence item we're working on
-        int seq_start = 0;
-        int b = 0;
-        while ((b < B) && (seq_start+lengths[b]) < row)
+        while ((b < B) && (seq_start+lengths[b]) < t)
         {
            seq_start += lengths[b];
            b += 1;
         }
             
-        dX = &dX[row * O];
-        d_sum = &d_sum[b * O];
+        float* dX_t = &dX[t * O];
+        const float* d_sum_b = &d_sum_b[b * O];
 
         for (int i=0; i < O; ++i) 
         {
-            dX[i] = d_sum[i];
-        }
-    }
-}
-
-
-extern "C" __global__
-void backprop_mean_pool(float* dX, const float* d_sum, const int* lengths,
+            dX_t[i] = d_sum_b[i];
+        }
+    }
+}
+
+
+extern "C" __global__
+void backprop_mean_pool(float* dX, const float* d_mean, 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)
+    int _loop_start = blockIdx.x * blockDim.x + threadIdx.x;
+    int _loop_stride = blockDim.x * gridDim.x;
+    int seq_start = 0;
+    int b = 0;
+    for (int t = _loop_start; t < T; t += _loop_stride)
     { 
         // Find the sequence item we're working on
-        int seq_start = 0;
-        int b = 0;
-        while ((b < B) && (seq_start+lengths[b]) < row)
+        while ((b < B) && (seq_start+lengths[b]) < t)
         {
            seq_start += lengths[b];
            b += 1;
         }
             
-        dX = &dX[row * O];
-        d_sum = &d_sum[b * O];
+        dX_t = &dX[t * O];
+        d_mean_b = &d_mean_b[b * O];
+        int lengths_b = lengths[b];
 
         for (int i=0; i < O; ++i) 
         {
-            dX[i] = d_sum[i] / lengths[i];
+            dX_t[i] = d_mean_t[i] / lengths_b;
         }
     }
 }

}



extern "C" __global__
void sum_pool(float* output,
    const float* X, const int* lengths, int B, int T, int O)
{
    // Compute sums of a batch of concatenated sequences
    for (int b = blockIdx.x * blockDim.x + threadIdx.x; b < B; b += blockDim.x * gridDim.x)
    {
        // Go to the regions we're working on
        for (int i=0; i < b; ++i) {

    // * x_start=-6, x_end=9 : (0-2) * 3, (0+2+1) * 3
    // * x_start=-3, x_end=13 : (1-2) * 3, (1+2+1) * 3
    // * x_start=0, x_end=16 : (2-2) * 3, (2+2+1) * 3
    int _loop_start = blockIdx.x * blockDim.x + threadIdx.x;
    int _loop_stride = blockDim.x * gridDim.x;
    int nF = nW * 2 + 1;
    for (int b = _loop_start; b < B; b += _loop_stride)
    {
        int o_start = b * I * nF;
        // Let's say b=0, nW=1, I=10, B=20
        // x_start = (0-1) * 10 : -10
