Update Dockerfile.sm

awslabs · Sep 30, 2024 · 9a48428 · 9a48428
1 parent d67e009
commit 9a48428
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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