diff --git a/batch_run.sh b/batch_run.sh
index b49ed766..304a7073 100644
--- a/batch_run.sh
+++ b/batch_run.sh
@@ -4,7 +4,7 @@
 
 repeat_cnt=$1
 
-export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
 
 for ((i=0; i<repeat_cnt; i++)); do
     bash run_once.sh lora_sft Qwen-7B 8 500
diff --git a/run_once.sh b/run_once.sh
index c7bca5ef..d41a3ba9 100644
--- a/run_once.sh
+++ b/run_once.sh
@@ -42,7 +42,7 @@ bash npu_status.sh ${output_dir} 60 0 &
 npu_status_pid=$!
 # echo "${npu_status_pid}"
 
-if [ "${gpu_cnt}"="1" ]; then
+if [ "${gpu_cnt}" = "1" ]; then
    ASCEND_RT_VISIBLE_DEVICES=0 llamafactory-cli train ${output_dir}/${run_name}.yaml | tee "${output_dir}/log.txt" &
    train_pid=$!
    echo "Start single npu train"