Dockerfile: Support CUDA 11 (pytorch#45071)

Summary:
Although PyTorch already supports CUDA 11, the Dockerfile still relies on CUDA 10. This pull request upgrades all the necessary versions such that recent NVIDIA GPUs like A100 can be used.

Pull Request resolved: pytorch#45071

Reviewed By: ezyang

Differential Revision: D23873224

Pulled By: seemethere

fbshipit-source-id: 822c25f183dcc3b4c5b780c00cd37744d34c6e00

Author: tindzk

Dockerfile

@@ -44,13 +44,13 @@ WORKDIR /opt/pytorch
 COPY --from=conda /opt/conda /opt/conda
 COPY --from=submodule-update /opt/pytorch /opt/pytorch
 RUN --mount=type=cache,target=/opt/ccache \
-    TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
+    TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX 8.0" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
     CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
     python setup.py install
 
 FROM conda as conda-installs
 ARG INSTALL_CHANNEL=pytorch-nightly
-RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -y pytorch torchvision cudatoolkit=10.1 && \
+RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -y pytorch torchvision cudatoolkit=11.0.221 && \
     /opt/conda/bin/conda clean -ya
 
 FROM ${BASE_IMAGE} as official

docker.Makefile

@@ -9,7 +9,7 @@ DOCKER_ORG       = $(shell whoami)
 endif
 
 BASE_RUNTIME     = ubuntu:18.04
-BASE_DEVEL       = nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04
+BASE_DEVEL       = nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04
 
 # The conda channel to use to install pytorch / torchvision
 INSTALL_CHANNEL  = pytorch