From b7dd39a6a747e19316959f41a90a6ea6b18289e4 Mon Sep 17 00:00:00 2001 From: jalencato Date: Mon, 30 Sep 2024 15:23:52 -0700 Subject: [PATCH] Update Dockerfile.sm --- docker/sagemaker/Dockerfile.sm | 2 ++ 1 file changed, 2 insertions(+) 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