From 0b7dc0192128f6d53b0f21c69ff7c324604a2628 Mon Sep 17 00:00:00 2001
From: SimonLee <lixiaomin_235@163.com>
Date: Sun, 18 Feb 2024 11:21:50 +0800
Subject: [PATCH] fix error when only 1 gpu

---
 ptuning/ds_train_finetune.sh | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/ptuning/ds_train_finetune.sh b/ptuning/ds_train_finetune.sh
index b9f06fe..c7e6baa 100644
--- a/ptuning/ds_train_finetune.sh
+++ b/ptuning/ds_train_finetune.sh
@@ -2,8 +2,9 @@
 LR=1e-4
 
 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 \
     --do_train \
     --train_file AdvertiseGen/train.json \