correcting cuda 11.2 image name in CI and CD (apache#19960)

Co-authored-by: Rohit Kumar Srivastava <[email protected]>
access2rohit · Mar 9, 2021 · e16818d · e16818d
1 parent 2e75bdd
commit e16818d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/cd/utils/mxnet_base_image.sh b/cd/utils/mxnet_base_image.sh
@@ -33,6 +33,9 @@ case ${mxnet_variant} in
     cu110*)
     echo "nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04"
     ;;
+    cu112*)
+    echo "nvidia/cuda:11.2.1-cudnn8-runtime-ubuntu18.04"
+    ;;
     cpu)
     echo "ubuntu:18.04"
     ;;

diff --git a/ci/docker/Dockerfile.build.ubuntu_blc b/ci/docker/Dockerfile.build.ubuntu_blc
@@ -18,7 +18,7 @@
 #
 # Dockerfile to build and run MXNet on Ubuntu 16.04 for CPU
 
-FROM ubuntu:16.04
+FROM nvidia/cuda:11.2.1-cudnn8-devel-ubuntu16.04
 
 WORKDIR /work/deps