diff --git a/test_tipc/common_func.sh b/test_tipc/common_func.sh
index 3f0fa66b77..85dfe21725 100644
--- a/test_tipc/common_func.sh
+++ b/test_tipc/common_func.sh
@@ -30,6 +30,7 @@ function func_set_params(){
 
 function func_parser_params(){
     strs=$1
+    MODE=$2
     IFS=":"
     array=(${strs})
     key=${array[0]}
diff --git a/test_tipc/test_train_inference_python.sh b/test_tipc/test_train_inference_python.sh
index 321eeb5534..6773c44e9e 100644
--- a/test_tipc/test_train_inference_python.sh
+++ b/test_tipc/test_train_inference_python.sh
@@ -20,10 +20,10 @@ train_use_gpu_value=$(func_parser_value "${lines[4]}")
 autocast_list=$(func_parser_value "${lines[5]}")
 autocast_key=$(func_parser_key "${lines[5]}")
 epoch_key=$(func_parser_key "${lines[6]}")
-epoch_num=$(func_parser_params "${lines[6]}")
+epoch_num=$(func_parser_params "${lines[6]}" "${MODE}")
 save_model_key=$(func_parser_key "${lines[7]}")
 train_batch_key=$(func_parser_key "${lines[8]}")
-train_batch_value=$(func_parser_params "${lines[8]}")
+train_batch_value=$(func_parser_params "${lines[8]}" "${MODE}")
 pretrain_model_key=$(func_parser_key "${lines[9]}")
 pretrain_model_value=$(func_parser_value "${lines[9]}")
 train_model_name=$(func_parser_value "${lines[10]}")