Merge branch 'main' into gsp_process_time_out

awslabs · Dec 13, 2024 · 52971ab · 52971ab
2 parents cb1a945 + 66b1998
commit 52971ab
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 0 deletions.
diff --git a/docker/local/Dockerfile.local b/docker/local/Dockerfile.local
@@ -60,6 +60,7 @@ RUN TORCH_MAJOR_MINOR=$(echo $TORCH_VERSION | cut -c1-3) && \
     pip install \
     dgl==${DGL_VERSION} -f https://data.dgl.ai/wheels/torch-${TORCH_MAJOR_MINOR}/repo.html \
     ogb==${OGB_VERSION} \
+    torchdata==0.9.0 \
     transformers==${TRANSFORMERS_VERSION} \
     && rm -rf /root/.cache
 
@@ -72,6 +73,7 @@ RUN TORCH_MAJOR_MINOR=$(echo $TORCH_VERSION | cut -c1-3) && \
     -f https://data.dgl.ai/wheels/torch-${TORCH_MAJOR_MINOR}/cu${DGL_CUDA_VERSION}/repo.html \
     ogb==${OGB_VERSION} \
     torch==${TORCH_VERSION} \
+    torchdata==0.9.0 \
     transformers==${TRANSFORMERS_VERSION} \
     && rm -rf /root/.cache
 

diff --git a/docker/parmetis/Dockerfile.parmetis b/docker/parmetis/Dockerfile.parmetis
@@ -40,6 +40,7 @@ RUN pip install \
     pydantic==2.7.0 \
     scikit-learn==1.4.2 \
     scipy==1.13.0 \
+    torchdata==0.9.0 \
     pyyaml \
     && rm -rf /root/.cache
 

diff --git a/docker/sagemaker/Dockerfile.sm b/docker/sagemaker/Dockerfile.sm
@@ -28,8 +28,10 @@ RUN pip3 install \
         numpy==1.26.4 \
         ogb==1.3.6 \
         pyarrow==16.1.0 \
+        pydantic==2.7.1 \
         scikit-learn==1.5.0 \
         scipy==1.13.1 \
+        torchdata==0.9.0 \
         transformers==4.28.1 \
     && rm -rf /root/.cache