Install PyTorch for cuda 9

songjiaru · Feb 20, 2019 · 9fa062c · 9fa062c
1 parent 02f8400
commit 9fa062c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gpu.Dockerfile b/gpu.Dockerfile
@@ -49,7 +49,7 @@ RUN pip uninstall -y tensorflow && \
     pip install /tmp/tensorflow_gpu/tensorflow*.whl && \
     rm -rf /tmp/tensorflow_gpu && \
     conda uninstall -y pytorch-cpu torchvision-cpu && \
-    conda install -y pytorch torchvision -c pytorch && \
+    conda install -y pytorch torchvision cudatoolkit=9.0 -c pytorch && \
     /tmp/clean-layer.sh
 
 # Install GPU-only packages