diff --git a/nlp/llm/mixtral/openrlhf/train_sft_mixtral_lora.sh b/nlp/llm/mixtral/openrlhf/train_sft_mixtral_lora.sh
index f702d4554067328179134302256e037a40bd65e6..e60ba4dd65d091ed20d5bad89d6cec71332db97b 100644
--- a/nlp/llm/mixtral/openrlhf/train_sft_mixtral_lora.sh
+++ b/nlp/llm/mixtral/openrlhf/train_sft_mixtral_lora.sh
@@ -22,7 +22,8 @@ openrlhf.cli.train_sft \
     --learning_rate 5e-6 \
     --lora_rank 64 \
     --lora_alpha 64 \
-    --aux_loss_coef 0.001
+    --aux_loss_coef 0.001 \
+    --gradient_checkpointing_use_reentrant
 EOF
 
 if [[ ${1} != "slurm" ]]; then