diff --git a/Dockerfile b/Dockerfile index 51fe7be4..37bcb35d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -29,14 +29,16 @@ LABEL "nemo.library"=${IMAGE_LABEL} WORKDIR /opt # Install the minimal libcu* libraries needed by NeMo Curator -RUN conda create -y --name curator -c conda-forge -c nvidia \ +ENV _CUDA_VER=${CUDA_VER} +RUN conda create -y --name curator -c nvidia/label/cuda-${_CUDA_VER} -c conda-forge \ python=3.10 \ cuda-cudart \ libcufft \ libcublas \ libcurand \ libcusparse \ - libcusolver && \ + libcusolver \ + cuda-nvvm && \ source activate curator && \ pip install --upgrade pytest pip