[Minor] add device constrain in script.

efeslab · Apr 8, 2024 · bd96c1e · bd96c1e
1 parent 0dcfda9
commit bd96c1e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/scripts/run_atom_ppl.sh b/scripts/run_atom_ppl.sh
@@ -21,7 +21,8 @@ resultFile=$dir/atom_llama_ppl.csv
 logFile=$dir/atom_llama_w${BIT}a${BIT}.log
 touch $logFile
 
-CUDA_VISIBLE_DEVICE=0 python ${dir}/model/main.py ${MODEL} ${CALIB_DATA} \
+export CUDA_VISIBLE_DEVICES=0 # Avoid causing not same devices error in "auto" device_map.
+python ${dir}/model/main.py ${MODEL} ${CALIB_DATA} \
     ${cmd_base} ${cmd_group} ${cmd_reorder} ${cmd_clip} ${cmd_adv} ${cmd_eval} \
     2>&1 | tee ${logFile}