From 00ac9774eb480e51bd2e604df648dfbb0f01f401 Mon Sep 17 00:00:00 2001
From: moto <855818+mthrok@users.noreply.github.com>
Date: Wed, 20 Jan 2021 23:37:00 +0000
Subject: [PATCH] Update GPU test image for development

---
 .circleci/config.yml                              |  2 +-
 .circleci/config.yml.in                           |  2 +-
 .circleci/unittest/linux/docker/build_and_push.sh | 11 +++++------
 3 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index c40054d0e3..062e6cccc4 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -430,7 +430,7 @@ jobs:
     resource_class: gpu.small
     environment:
       <<: *environment
-      image_name: "pytorch/torchaudio_unittest_base:manylinux-cuda10.1"
+      image_name: pytorch/torchaudio_unittest_base:manylinux-cuda10.1-cudnn7-20201203
     steps:
       - checkout
       - attach_workspace:
diff --git a/.circleci/config.yml.in b/.circleci/config.yml.in
index 535729d72e..7258487862 100644
--- a/.circleci/config.yml.in
+++ b/.circleci/config.yml.in
@@ -430,7 +430,7 @@ jobs:
     resource_class: gpu.small
     environment:
       <<: *environment
-      image_name: "pytorch/torchaudio_unittest_base:manylinux-cuda10.1"
+      image_name: pytorch/torchaudio_unittest_base:manylinux-cuda10.1-cudnn7-20201203
     steps:
       - checkout
       - attach_workspace:
diff --git a/.circleci/unittest/linux/docker/build_and_push.sh b/.circleci/unittest/linux/docker/build_and_push.sh
index 8bf33eaaa3..e7ced13ad3 100755
--- a/.circleci/unittest/linux/docker/build_and_push.sh
+++ b/.circleci/unittest/linux/docker/build_and_push.sh
@@ -7,15 +7,14 @@ if [ $# -ne 1 ]; then
     exit 1
 fi
 
+datestr="$(date "+%Y%m%d")"
 if [ "$1" = "cpu" ]; then
     base_image="ubuntu:18.04"
-    image="pytorch/torchaudio_unittest_base:manylinux"
-elif [[ "$1" =~ ^(9.2|10.1)$ ]]; then
-    base_image="nvidia/cuda:$1-runtime-ubuntu18.04"
-    image="pytorch/torchaudio_unittest_base:manylinux-cuda$1"
+    image="pytorch/torchaudio_unittest_base:manylinux-${datestr}"
 else
-    printf "Unexpected <CUDA_VERSION> string: %s" "$1"
-    exit 1;
+    base_image="nvidia/cuda:$1-devel-ubuntu18.04"
+    docker pull "${base_image}"
+    image="pytorch/torchaudio_unittest_base:manylinux-cuda$1-${datestr}"
 fi
 
 cd "$( dirname "${BASH_SOURCE[0]}" )"