diff --git a/docker/sagemaker/Dockerfile.sm b/docker/sagemaker/Dockerfile.sm
index db760061b5..4620e3e079 100644
--- a/docker/sagemaker/Dockerfile.sm
+++ b/docker/sagemaker/Dockerfile.sm
@@ -6,7 +6,7 @@ ARG DEVICE=gpu
 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
 # Install DGL GPU version
-RUN pip3 install dgl==2.3.0 -f https://data.dgl.ai/wheels/torch-2.3/repo.html && rm -rf /root/.cache
+RUN pip3 install dgl==2.3.0 -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