diff --git a/b_asic/codegen/testbench/streaming_matrix_transposition_tb.vhdl b/b_asic/codegen/testbench/streaming_matrix_transposition_tb.vhdl
index 398a50073cc6c7e7349f45d726d63b0d0cb93c5c..e91a898dfe02183bcb32318e32f8a92c1da65d8f 100644
--- a/b_asic/codegen/testbench/streaming_matrix_transposition_tb.vhdl
+++ b/b_asic/codegen/testbench/streaming_matrix_transposition_tb.vhdl
@@ -40,7 +40,7 @@ begin
     -- Input generation
     input_gen_proc: process begin
         wait until en = '1';
-        for i in 0 to ROWS*COLS-1 loop
+        for i in 0 to 4*ROWS*COLS-1 loop
             wait until clk = '0';
             input <= std_logic_vector(to_unsigned(i, input'length));
         end loop;
@@ -51,10 +51,17 @@ begin
     output_test_proc: process begin
         wait until en = '1';
         wait until output = std_logic_vector(to_unsigned(0, output'length));
-        for col in 0 to COLS-1 loop
-            for row in 0 to ROWS-1 loop
-                wait until clk = '0';
-                --check(output = std_logic_vector(to_unsigned(row*COLS + col, output'length)));
+        for i in 0 to 3 loop
+            for col in 0 to COLS-1 loop
+                for row in 0 to ROWS-1 loop
+                    wait until clk = '0';
+                    check(
+                        output =
+                        std_logic_vector(
+                            to_unsigned(i*ROWS*COLS + row*COLS + col, output'length)
+                        )
+                    );
+                end loop;
             end loop;
         end loop;
         done <= true;