diff --git a/scripts/deepspeed_gpt3_small.sh b/scripts/deepspeed_gpt3_small.sh
index f0185b5..516ffbc 100644
--- a/scripts/deepspeed_gpt3_small.sh
+++ b/scripts/deepspeed_gpt3_small.sh
@@ -39,7 +39,7 @@ gpt_options=" \
        --deepspeed_config src/deepspeed_config/gpt3_small_2048.json \
 "
 
-run_cmd="mpirun --np ${NUM_GPUS_PER_WORKER} python ../pretrain_gpt3.py $@ ${gpt_options}"
+run_cmd="USE_DEEPSPEED=1 mpirun --np ${NUM_GPUS_PER_WORKER} python ../pretrain_gpt3.py $@ ${gpt_options}"
 echo "${run_cmd}"
 eval "${run_cmd}"