Update building docker images (#12837)

Co-authored-by: Akihiro Nitta <akihiro@pytorchlightning.ai>

.github/workflows/ci_dockers.yml

@@ -73,11 +73,14 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python_version: ["3.7", "3.9"]
-        pytorch_version: ["1.8", "1.11"]
         include:
           # the config used in '.azure-pipelines/gpu-tests.yml'
-          - {python_version: "3.9", pytorch_version: "1.10"}
+          - {python_version: "3.7", pytorch_version: "1.10", cuda_version: "11.1"}
+          - {python_version: "3.7", pytorch_version: "1.11", cuda_version: "11.3.1"}
+          # latest (used in Tutorials)
+          - {python_version: "3.8", pytorch_version: "1.8", cuda_version: "11.1"}
+          - {python_version: "3.9", pytorch_version: "1.10", cuda_version: "11.1"}
+          - {python_version: "3.9", pytorch_version: "1.11", cuda_version: "11.3.1"}
     steps:
       - name: Checkout
         uses: actions/checkout@v2
@@ -88,6 +91,7 @@ jobs:
           build-args: |
             PYTHON_VERSION=${{ matrix.python_version }}
             PYTORCH_VERSION=${{ matrix.pytorch_version }}
+            CUDA_VERSION=${{ matrix.cuda_version }}
           file: dockers/base-cuda/Dockerfile
           push: false
         timeout-minutes: 75

.github/workflows/events-nightly.yml

@@ -115,12 +115,12 @@ jobs:
       matrix:
         include:
           # the config used in '.azure-pipelines/gpu-tests.yml'
-          - {python_version: "3.7", pytorch_version: "1.8"}
-          - {python_version: "3.7", pytorch_version: "1.10"}
+          - {python_version: "3.7", pytorch_version: "1.10", cuda_version: "11.1"}
+          - {python_version: "3.7", pytorch_version: "1.11", cuda_version: "11.3.1"}
           # latest (used in Tutorials)
-          - {python_version: "3.8", pytorch_version: "1.8"}
-          - {python_version: "3.9", pytorch_version: "1.10"}
-          - {python_version: "3.9", pytorch_version: "1.11"}
+          - {python_version: "3.8", pytorch_version: "1.8", cuda_version: "11.1"}
+          - {python_version: "3.9", pytorch_version: "1.10", cuda_version: "11.1"}
+          - {python_version: "3.9", pytorch_version: "1.11", cuda_version: "11.3.1"}
 
     steps:
       - name: Checkout
@@ -140,6 +140,7 @@ jobs:
           build-args: |
             PYTHON_VERSION=${{ matrix.python_version }}
             PYTORCH_VERSION=${{ matrix.pytorch_version }}
+            CUDA_VERSION=${{ matrix.cuda_version }}
           file: dockers/base-cuda/Dockerfile
           push: ${{ env.PUSH_TO_HUB }}
           tags: pytorchlightning/pytorch_lightning:base-cuda-py${{ matrix.python_version }}-torch${{ matrix.pytorch_version }}

dockers/base-cuda/Dockerfile

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG CUDA_VERSION=11.1
+ARG CUDA_VERSION=11.3.1
 
 FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04