diff --git a/crates/burn-jit/src/kernel/matmul/tiling2d_cube/load_shared_memory.rs b/crates/burn-jit/src/kernel/matmul/tiling2d_cube/load_shared_memory.rs
index 4bd015c6b..bbb2ed0d5 100644
--- a/crates/burn-jit/src/kernel/matmul/tiling2d_cube/load_shared_memory.rs
+++ b/crates/burn-jit/src/kernel/matmul/tiling2d_cube/load_shared_memory.rs
@@ -200,7 +200,7 @@ fn write_tile_transposed<F: Float>(
     } else {
         for i in range(0u32, Comptime::get(tile_size), unroll) {
             let mut transposed = F::vectorized(0., Comptime::get(tile_size));
-            for j in range(0u32, Comptime::get(tile_size), unroll) {
+            for j in range(0u32, Comptime::get(tile_size), Comptime::new(true)) {
                 let mut row = tile[j];
                 if j > Comptime::runtime(tile_size) {
                     row = F::vectorized(0., Comptime::get(tile_size));