Use CUDA 11.2 for nightly docker build. (#51990)

Summary:
Set CUDA_VERSION to 11.2.0 since Nvidia name their docker image on Ubuntu 18.04 to be nvidia/cuda:11.2.0-cudnn8-devel-ubuntu18.04.

Note that cudatoolkit 11.2.0 is not yet on [conda](https://repo.anaconda.com/pkgs/main/linux-64/), and we need to wait for that before merging this PR.

- https://hub.docker.com/r/nvidia/cuda/

Pull Request resolved: https://github.com/pytorch/pytorch/pull/51990

Reviewed By: samestep

Differential Revision: D26371193

Pulled By: xuzhao9

fbshipit-source-id: 76915490dc30ddb03ceeeadb3c45a6c02b60401e
This commit is contained in:
Xu Zhao
2021-02-10 10:43:22 -08:00
committed by Facebook GitHub Bot
parent c4a8f0ceaa
commit 03e82f7944
3 changed files with 7 additions and 1 deletions

View File

@ -52,9 +52,10 @@ RUN --mount=type=cache,target=/opt/ccache \
FROM conda as conda-installs
ARG PYTHON_VERSION=3.8
ARG CUDA_VERSION=11.0
ARG CUDA_CHANNEL=defaults
ARG INSTALL_CHANNEL=pytorch-nightly
ENV CONDA_OVERRIDE_CUDA=${CUDA_VERSION}
RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -y python=${PYTHON_VERSION} pytorch torchvision torchtext "cudatoolkit=${CUDA_VERSION}" && \
RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -c "${CUDA_CHANNEL}" -y python=${PYTHON_VERSION} pytorch torchvision torchtext "cudatoolkit=${CUDA_VERSION}" && \
/opt/conda/bin/conda clean -ya
RUN /opt/conda/bin/pip install torchelastic