diff --git a/dev/cuda/attention_backward.cu b/dev/cuda/attention_backward.cu
index 7a53a22..3d57c8f 100644
--- a/dev/cuda/attention_backward.cu
+++ b/dev/cuda/attention_backward.cu
@@ -633,8 +633,8 @@ __global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_ker
 }
 
 template<int BlockSize>
-__global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_kernel7(float* dpreatt, const float* datt, const float* att,
-                                                                                     int B, int T, int C, int NH) {
+__global__ void softmax_autoregressive_backward_kernel7(float* dpreatt, const float* datt, const float* att,
+                                                        int B, int T, int C, int NH) {
     namespace cg = cooperative_groups;
     cg::thread_block block = cg::this_thread_block();
     cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
@@ -657,29 +657,21 @@ __global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_ker
         block_acc[warp.thread_rank()] = 0;
     }
 
-    int block_steps = ceil_div(t+1, BlockSize);
-    // very important: This loop condition needs to be the same for all threads.
-    // even if a thread later on is not going to do any work, it needs to participate in the
-    // data loading process!
-    for (int t3f = 0; t3f < block_steps; ++t3f) {
-        int t3 = t3f * BlockSize + block.thread_rank();
+    float local_sum = 0;
+    for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) {
+        local_sum += att_bth[t2] * datt_bth[t2];
+    }
 
-        float at3 = att_bth[min(t, t3)];
-        float local_sum = 0;
-        for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) {
-            local_sum += att_bth[t2] * datt_bth[t2];
-        }
-        block.sync();
-        block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus<float>{});
-        block.sync();
-        local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
+    block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus<float>{});
+    block.sync();
+    local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
 
+    for (int t3 = block.thread_rank(); t3 <= t; t3  += BlockSize) {
+        float at3 = att_bth[t3];
         float acc = -local_sum * at3;
-        float at_t2_eq_t3 = at3 * datt_bth[min(t, t3)];
+        float at_t2_eq_t3 = at3 * datt_bth[t3];
         acc += (at_t2_eq_t3 * (1.f - at3) - at_t2_eq_t3 * (0.f - at3));
-        if(t3 <= t) {
-            dpreatt_bth[t3] = scale * acc;
-        }
+        dpreatt_bth[t3] = scale * acc;
     }
 }
 
@@ -930,7 +922,7 @@ int main(int argc, char **argv) {
     srand(0); // reproducibility
 
     // hyperparameters
-    int B = 8;
+    int B = 4;
     int T = 1024;
     int C = 768;
     int NH = 12;
diff --git a/train_gpt2.cu b/train_gpt2.cu
index c7a6d0c..1ec64d2 100644
--- a/train_gpt2.cu
+++ b/train_gpt2.cu
@@ -690,29 +690,21 @@ __global__ void softmax_autoregressive_backward_kernel(float* dpreatt, const flo
         block_acc[warp.thread_rank()] = 0;
     }
 
-    int block_steps = CEIL_DIV(t+1, BlockSize);
-    // very important: This loop condition needs to be the same for all threads.
-    // even if a thread later on is not going to do any work, it needs to participate in the
-    // data loading process!
-    for (int t3f = 0; t3f < block_steps; ++t3f) {
-        int t3 = t3f * BlockSize + block.thread_rank();
+    float local_sum = 0;
+    for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) {
+        local_sum += att_bth[t2] * datt_bth[t2];
+    }
 
-        float at3 = att_bth[min(t, t3)];
-        float local_sum = 0;
-        for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) {
-            local_sum += att_bth[t2] * datt_bth[t2];
-        }
-        block.sync();
-        block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus<float>{});
-        block.sync();
-        local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
+    block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus<float>{});
+    block.sync();
+    local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
 
+    for (int t3 = block.thread_rank(); t3 <= t; t3  += BlockSize) {
+        float at3 = att_bth[t3];
         float acc = -local_sum * at3;
-        float at_t2_eq_t3 = at3 * datt_bth[min(t, t3)];
+        float at_t2_eq_t3 = at3 * datt_bth[t3];
         acc += (at_t2_eq_t3 * (1.f - at3) - at_t2_eq_t3 * (0.f - at3));
-        if(t3 <= t) {
-            dpreatt_bth[t3] = scale * acc;
-        }
+        dpreatt_bth[t3] = scale * acc;
     }
 }