diff --git a/docker/Dockerfile b/docker/Dockerfile index 7d4ae30746dfcf5e2104f9e3ea611812cf4be0ec..4710f730ebfab5f6eeece8a986529e436d273fc5 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,9 +1,13 @@ -ARG PYTORCH="1.1.0" -ARG CUDA="10.0" -ARG CUDNN="7.5" +ARG PYTORCH="1.3" +ARG CUDA="10.1" +ARG CUDNN="7" FROM pytorch/pytorch:${PYTORCH}-cuda${CUDA}-cudnn${CUDNN}-devel +ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0+PTX" +ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all" +ENV CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" + RUN apt-get update && apt-get install -y libglib2.0-0 libsm6 libxrender-dev libxext6 \ && apt-get clean \ && rm -rf /var/lib/apt/lists/*