diff --git a/finetune/install-wsl-dep-and-train.sh b/finetune/install-wsl-dep-and-train.sh
index 461edb3..4a2d25f 100644
--- a/finetune/install-wsl-dep-and-train.sh
+++ b/finetune/install-wsl-dep-and-train.sh
@@ -50,8 +50,9 @@ echo "loading $loadModel"
 modelInfo=$(python3 ./finetune/get_layer_and_embd.py $loadModel 5.2)
 echo $modelInfo
 if [[ $modelInfo =~ "--n_layer" ]]; then
+  sudo rm -rf /root/.cache/torch_extensions
   python3 ./finetune/lora/$modelInfo $@ --proj_dir lora-models --data_type binidx --lora \
-    --lora_parts=att,ffn,time,ln --strategy deepspeed_stage_2 --accelerator gpu
+    --lora_parts=att,ffn,time,ln --strategy deepspeed_stage_2 --accelerator gpu --ds_bucket_mb 2
 else
   echo "modelInfo is invalid"
   exit 1
diff --git a/finetune/requirements.txt b/finetune/requirements.txt
index 2e218e8..09c962b 100644
--- a/finetune/requirements.txt
+++ b/finetune/requirements.txt
@@ -1,3 +1,3 @@
-torch==1.13.1
+torch==2.1.2
 pytorch_lightning==1.9.5
-deepspeed==0.11.2
+deepspeed==0.12.6