Switch to using cudnn base image.

PiperOrigin-RevId: 387058301 Change-Id: If7d2716bfe89228e9fb8f909b72c425bd89f1f7f

Switch to using cudnn base image.
PiperOrigin-RevId: 387058301 Change-Id: If7d2716bfe89228e9fb8f909b72c425bd89f1f7f
acf25fc8 · DeepMind · Copybara-Service · 2028cd0a · acf25fc8
Commit acf25fc8 authored Jul 27, 2021 by DeepMind Committed by Copybara-Service Jul 27, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

docker/Dockerfile docker/Dockerfile +2 -2

No files found.
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -13,7 +13,7 @@
 # limitations under the License.

 ARG CUDA=11.0
-FROM nvidia/cuda:${CUDA}-base
+FROM nvidia/cuda:${CUDA}-cudnn8-runtime-ubuntu16.04
 # FROM directive resets ARGS, so we specify again (the value is retained if
 # previously set).
 ARG CUDA
@@ -52,7 +52,7 @@ ENV PATH="/opt/conda/bin:$PATH"
 RUN conda update -qy conda \
    && conda install -y -c conda-forge \
      openmm=7.5.1 \
-      cudatoolkit==${CUDA}.3 \
+      cudatoolkit==${CUDA_VERSION} \
      pdbfixer \
      pip \
      python=3.7