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

.github/scripts/build_publish_nightly_docker.sh

@@ -9,6 +9,8 @@ PYTORCH_DOCKER_TAG=$(git describe --tags --always)-devel
 make -f docker.Makefile \
      DOCKER_REGISTRY=ghcr.io \
      DOCKER_ORG=pytorch \
+     CUDA_VERSION=11.2.0 \
+     CUDA_CHANNEL=conda-forge \
      DOCKER_IMAGE=pytorch-nightly \
      DOCKER_TAG=${PYTORCH_DOCKER_TAG} \
      INSTALL_CHANNEL=pytorch-nightly BUILD_TYPE=official devel-image

Dockerfile

@@ -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
 

docker.Makefile

@@ -13,6 +13,8 @@ CUDNN_VERSION             = 8
 BASE_RUNTIME              = ubuntu:18.04
 BASE_DEVEL                = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_VERSION)-devel-ubuntu18.04
 
+# The conda channel to use to install cudatoolkit
+CUDA_CHANNEL              = defaults
 # The conda channel to use to install pytorch / torchvision
 INSTALL_CHANNEL           = pytorch
 
@@ -24,6 +26,7 @@ BUILD_PROGRESS            = auto
 BUILD_ARGS                = --build-arg BASE_IMAGE=$(BASE_IMAGE) \
 							--build-arg PYTHON_VERSION=$(PYTHON_VERSION) \
 							--build-arg CUDA_VERSION=$(CUDA_VERSION) \
+							--build-arg CUDA_CHANNEL=$(CUDA_CHANNEL) \
 							--build-arg PYTORCH_VERSION=$(PYTORCH_VERSION) \
 							--build-arg INSTALL_CHANNEL=$(INSTALL_CHANNEL)
 EXTRA_DOCKER_BUILD_FLAGS ?=