update pytorch docker to cuda 10

Summary: Pull Request resolved: pytorch/pytorch#16082 Differential Revision: D13699081 Pulled By: soumith fbshipit-source-id: 86942e2c5595931384cf87dd1ef75936a4d74a57
TonyTangYu · Jan 16, 2019 · aff0964 · aff0964
1 parent d33e7d1
commit aff0964
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/docker/pytorch/Dockerfile b/docker/pytorch/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04
+FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu16.04
 ARG PYTHON_VERSION=3.6
 RUN apt-get update && apt-get install -y --no-install-recommends \
          build-essential \
@@ -17,9 +17,10 @@ RUN curl -o ~/miniconda.sh -O  https://repo.continuum.io/miniconda/Miniconda3-la
      ~/miniconda.sh -b -p /opt/conda && \
      rm ~/miniconda.sh && \
      /opt/conda/bin/conda install -y python=$PYTHON_VERSION numpy pyyaml scipy ipython mkl mkl-include cython typing && \
-     /opt/conda/bin/conda install -y -c pytorch magma-cuda90 && \
+     /opt/conda/bin/conda install -y -c pytorch magma-cuda100 && \
      /opt/conda/bin/conda clean -ya
 ENV PATH /opt/conda/bin:$PATH
+RUN pip install ninja
 # This must be done before pip so that requirements.txt is available
 WORKDIR /opt/pytorch
 COPY . .