fix error when only 1 gpu

pull/1455/head
SimonLee 2024-02-18 11:21:50 +08:00
parent a0dc8f774a
commit 0b7dc01921
1 changed files with 2 additions and 1 deletions

View File

@ -2,8 +2,9 @@
LR=1e-4 LR=1e-4
MASTER_PORT=$(shuf -n 1 -i 10000-65535) MASTER_PORT=$(shuf -n 1 -i 10000-65535)
NUM_GPUS=$(nvidia-smi --list-gpus | wc -l)
deepspeed --num_gpus=1 --master_port $MASTER_PORT main.py \ deepspeed --num_gpus=$NUM_GPUS --master_port $MASTER_PORT main.py \
--deepspeed deepspeed.json \ --deepspeed deepspeed.json \
--do_train \ --do_train \
--train_file AdvertiseGen/train.json \ --train_file AdvertiseGen/train.json \