diff --git a/docker/sagemaker/Dockerfile.sm b/docker/sagemaker/Dockerfile.sm
index 7a0bd7067..97df1ba22 100644
--- a/docker/sagemaker/Dockerfile.sm
+++ b/docker/sagemaker/Dockerfile.sm
@@ -6,11 +6,13 @@ ARG DGL_VERSION=2.3.0
 
 FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.3.0-gpu-py311-cu121-ubuntu20.04-sagemaker as branch-gpu
 ENV dev_type=GPU
+ARG DGL_VERSION
 # Install DGL GPU version
 RUN pip3 install dgl==${DGL_VERSION}+cu121 -f https://data.dgl.ai/wheels/torch-2.3/cu121/repo.html && rm -rf /root/.cache
 
 FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.3.0-cpu-py311-ubuntu20.04-sagemaker as branch-cpu
 ENV dev_type=CPU
+ARG DGL_VERSION
 # Install DGL CPU version
 RUN pip3 install dgl==${DGL_VERSION} -f https://data.dgl.ai/wheels/torch-2.3/repo.html && rm -rf /root/.cache