     for (int t = _loop_start; t < T; t += _loop_stride)
     {
         // Find the sequence item we're working on
-        while ((b < B) && (seq_start+lengths[b]) < t)
+        while ((b < B) && (seq_start+lengths[b]) <= t)
         {
            seq_start += lengths[b];
            b += 1;

    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* dX_b = &dX[b*O*P];
	const float* dY_b = &dY[b*O];
	const int* which_b = &which[b*O];
        for (int i=0; i < O; ++i)
            dX_b[(i*P)+which_b[i]] = dY_b[i];
    }

    for (int b = _loop_start; b < B; b += _loop_stride)
    {
        // Go to the regions we're working on
	    float* output_b = &output[b*O];
        // Find the sequence item we're working on
	    int t = 0;
        for (int i=0; i < b; ++i) {
	        t += lengths[i];
        }
        int length = lengths[b];
        // Each invocation of the kernel sums one batch.
        for (int i=0; i < length; ++i) // Iterate over rows
        {
	        const float* X_t = &X[(t+i)*O];
            for (int j=0; j < O; ++j)
            {
              output_b[j] += X_t[j];
