devel

HumanSignal · Dec 28, 2024 · 923925e · 923925e
1 parent f896e86
commit 923925e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/label_studio_ml/examples/grounding_dino/Dockerfile b/label_studio_ml/examples/grounding_dino/Dockerfile
@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
+FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-devel
 ARG DEBIAN_FRONTEND=noninteractive
 ARG TEST_ENV
 

diff --git a/label_studio_ml/examples/grounding_sam/Dockerfile b/label_studio_ml/examples/grounding_sam/Dockerfile
@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
+FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-devel
 ARG DEBIAN_FRONTEND=noninteractive
 ARG TEST_ENV
 
@@ -24,6 +24,7 @@ ENV PYTHONUNBUFFERED=1 \
     CUDA_HOME=/usr/local/cuda \
     GROUNDINGDINO_REPO_PATH=/GroundingDINO
 
+RUN conda install -c "nvidia/label/cuda-12.1.1" cuda -y
 ENV CUDA_HOME=/opt/conda \
     TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6+PTX;8.9;9.0"