From 368a0c06d4ed9fccb1e03cfa570c4d8530be9203 Mon Sep 17 00:00:00 2001
From: atalman <atalman@fb.com>
Date: Tue, 19 Dec 2023 00:58:15 +0000
Subject: [PATCH] [releng] Docker Official release make sure cuda version is
 part of image name (#116070)

Follow up on https://github.com/pytorch/pytorch/pull/115949

Change docker build image name:
``pytorch:2.1.2-devel``-> ``2.1.2-cuda12.1-cudnn8-devel and 2.1.2-cuda11.8-cudnn8-devel``

Ref: https://github.com/orgs/pytorch/packages/container/package/pytorch-nightly

Naming will be same as in https://hub.docker.com/r/pytorch/pytorch/tags
Pull Request resolved: https://github.com/pytorch/pytorch/pull/116070
Approved by: https://github.com/huydhn, https://github.com/seemethere
---
 .github/workflows/docker-release.yml | 1 +
 docker.Makefile                      | 7 ++++---
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/docker-release.yml b/.github/workflows/docker-release.yml
index c678fccb104..a49a5189cbd 100644
--- a/.github/workflows/docker-release.yml
+++ b/.github/workflows/docker-release.yml
@@ -60,6 +60,7 @@ jobs:
       BUILD_IMAGE_TYPE: ${{ matrix.image_type }}
       BUILD_PLATFORMS: ${{ matrix.platform }}
       CUDA_VERSION: ${{ matrix.cuda_full_version }}
+      CUDA_VERSION_SHORT: ${{ matrix.cuda }}
       CUDNN_VERSION: ${{ matrix.cudnn_version }}
     steps:
       - name: Setup SSH (Click me for login details)
diff --git a/docker.Makefile b/docker.Makefile
index 63d4126f0df..1053f2c91fc 100644
--- a/docker.Makefile
+++ b/docker.Makefile
@@ -1,15 +1,16 @@
 DOCKER_REGISTRY          ?= docker.io
 DOCKER_ORG               ?= $(shell docker info 2>/dev/null | sed '/Username:/!d;s/.* //')
 DOCKER_IMAGE             ?= pytorch
-DOCKER_FULL_NAME          = $(DOCKER_REGISTRY)/$(DOCKER_ORG)/$(DOCKER_IMAGE)
+CUDA_VERSION_SHORT       ?= 12.1
+CUDA_VERSION             ?= 12.1.1
+CUDNN_VERSION            ?= 8
+DOCKER_FULL_NAME          = $(DOCKER_REGISTRY)/$(DOCKER_ORG)/$(DOCKER_IMAGE)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)
 
 ifeq ("$(DOCKER_ORG)","")
 $(warning WARNING: No docker user found using results from whoami)
 DOCKER_ORG                = $(shell whoami)
 endif
 
-CUDA_VERSION             ?= 12.1.1
-CUDNN_VERSION            ?= 8
 BASE_RUNTIME              = ubuntu:22.04
 BASE_DEVEL                = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_VERSION)-devel-ubuntu22.04
 CMAKE_VARS               ?=