diff --git a/Makefile b/Makefile
index 0303bd2..3b00113 100644
--- a/Makefile
+++ b/Makefile
@@ -17,3 +17,4 @@ clean:
 	rm -rf tb/test_source_code/**/*.bin
 	rm -rf tb/test_source_code/**/*.elf
 	rm -rf tb/test_source_code/**/*.o
+	rm -rf tb/test_source_code/**/*.tmp
diff --git a/scripts/gen_test.py b/scripts/gen_test.py
index d578d28..d68e4ce 100644
--- a/scripts/gen_test.py
+++ b/scripts/gen_test.py
@@ -1,6 +1,11 @@
 import re
+import sys
 
-source_code = open('test.S', 'r')
+if len(sys.argv) != 3:
+    print("Usage: python3 gen_test.py <path> <filename>")
+    exit(1)
+
+source_code = open(sys.argv[1] + "/" + sys.argv[2], 'r')
 Lines = source_code.readlines()
 test_file = []
 
@@ -47,7 +52,7 @@ for line in Lines:
 
 
 # save test_file to a file named test.tmp
-with open('test.tmp', 'w') as f:
+with open(sys.argv[1] + '/test.tmp', 'w') as f:
     for item in test_file:
         f.write("%s\n" % item)
 
@@ -66,6 +71,6 @@ for line in reversed(Lines):
         break
 
 # save test_file to a file named test.tmp
-with open('test.final.tmp', 'w') as f:
+with open(sys.argv[1] + '/test.final.tmp', 'w') as f:
     for item in final_test_file:
         f.write("%s\n" % item)
diff --git a/sim/Makefile b/sim/Makefile
index 1fdd80e..df8dfd8 100644
--- a/sim/Makefile
+++ b/sim/Makefile
@@ -1,5 +1,6 @@
 all:
 	./../scripts/get_bin.sh ../tb/test_source_code/tb_riscv_cpu/test
+	python3 ./../scripts/gen_test.py ../tb/test_source_code/tb_riscv_cpu test.S
 	vsim -c -do "do simu.do; quit -f"
 
 debug:
diff --git a/tb/tb_risc_v_cpu-dyn.v b/tb/tb_risc_v_cpu-dyn.v
index c6fe5d3..28b9a5f 100644
--- a/tb/tb_risc_v_cpu-dyn.v
+++ b/tb/tb_risc_v_cpu-dyn.v
@@ -53,7 +53,6 @@ module tb_risc_v_cpu ();
             read_instruction_2 = $fgetc(bin_file_inputs);
             read_instruction_3 = $fgetc(bin_file_inputs);
             read_instruction_4 = $fgetc(bin_file_inputs);
-            $display("read_instruction_1: %b", read_instruction_1);
 
             if (
                 read_instruction_1[8] != 1'b1 &&
@@ -91,8 +90,6 @@ module tb_risc_v_cpu ();
                 $finish;
             end
 
-            $display ("Line %d: %d:%b=%b ]", res, instruction_addr, reg_number, reg_test_value);
-
             if (test[instruction_addr][5:0] == 6'b111111) begin
                 test[instruction_addr][5:0] = reg_number;
                 test[instruction_addr][37:6] = reg_test_value;
@@ -156,7 +153,6 @@ module tb_risc_v_cpu ();
         while (!$feof(code_file_inputs))
         begin
             res = $fscanf(code_file_inputs, "%d=%d\n", reg_number, reg_test_value);
-            $display ("Line %d: %b=%b ]", res, reg_number, reg_test_value);
             if (res != 2) begin     // If fscanf failed, the test file structure is wrong, then exit
                 $display("Parsing test file failed");
                 $finish;
diff --git a/tb/test_source_code/tb_riscv_cpu/test.S b/tb/test_source_code/tb_riscv_cpu/test.S
index e6a95ce..bf7c7e1 100644
--- a/tb/test_source_code/tb_riscv_cpu/test.S
+++ b/tb/test_source_code/tb_riscv_cpu/test.S
@@ -1,11 +1,10 @@
 # t0 = 0
-li      t0, 0
-li      t2, 10 # R[2]=10, MEM[1]=6
+li      t0, 0  # R[5]=0
+li      t2, 10 # R[7]=10
 loop_head:
 bge     t0, t2, loop_end
 # Repeated code goes here
 addi    t0, t0, 1 # PC=16
 j		loop_head
 loop_end:
-
 # R[0]=0