Update Dockerfile.sm

awslabs · Sep 30, 2024 · b7dd39a · b7dd39a
1 parent d2ebac9
commit b7dd39a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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