conda: Add CUDA_HOME, cuda binaries to path (#1224)

pytorch · Dec 7, 2022 · 2946b23 · 2946b23
1 parent 5770ed3
commit 2946b23
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/conda/Dockerfile b/conda/Dockerfile
@@ -41,8 +41,12 @@ RUN bash ./install_conda.sh && rm install_conda.sh
 
 # Install CUDA
 FROM base as cuda
+ARG CUDA_VERSION=10.2
 RUN rm -rf /usr/local/cuda-*
 ADD ./common/install_cuda.sh install_cuda.sh
+ENV CUDA_HOME=/usr/local/cuda-${CUDA_VERSION}
+# Make things in our path by default
+ENV PATH=/usr/local/cuda-${CUDA_VERSION}/bin:$PATH
 
 FROM cuda as cuda11.6
 RUN bash ./install_cuda.sh 11.6

diff --git a/conda/build_docker.sh b/conda/build_docker.sh
@@ -52,6 +52,11 @@ if [[ "${DOCKER_TAG}" =~ ^cuda* ]]; then
     set -x
     docker tag ${DOCKER_IMAGE} "pytorch/conda-builder:cuda${CUDA_VERSION/./}"
   )
+  # Test that we're using the right CUDA compiler
+  (
+    set -x
+    docker run --rm "${DOCKER_IMAGE}" nvcc --version | grep "cuda_${CUDA_VERSION}"
+  )
 fi
 
 if [[ -n ${GITHUB_REF} ]]; then