change cuda version arg name
diff --git a/cd/python/docker/Dockerfile b/cd/python/docker/Dockerfile
index 7743c8b..cda17d3 100644
--- a/cd/python/docker/Dockerfile
+++ b/cd/python/docker/Dockerfile
@@ -23,11 +23,11 @@
ARG BASE_IMAGE
FROM ${BASE_IMAGE}
-ARG CUDA_VERSION
+ARG CUDA_VARIANT
# manually update cuda repo key: https://developer.nvidia.com/blog/updating-the-cuda-linux-gpg-repository-key/
# remove this after nvidia updates the base images with new keys
-RUN if [ "$CUDA_VERSION" != "false" ] ; then \
+RUN if [ "$CUDA_VARIANT" != "false" ] ; then \
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub; \
fi
@@ -49,9 +49,9 @@
tar -xvf arm-performance-libraries_21.0_Ubuntu-18.04_gcc-8.2.tar && \
arm-performance-libraries_21.0_Ubuntu-18.04_gcc-8.2/arm-performance-libraries_21.0_Ubuntu-18.04.sh -a; \
fi
-ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/opt/arm/armpl_21.0_gcc-8.2/lib:/usr/local/cuda-${CUDA_VERSION}/compat
+ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/opt/arm/armpl_21.0_gcc-8.2/lib:/usr/local/cuda-${CUDA_VARIANT}/compat
-RUN echo "CUDA_VERSION: ${CUDA_VERSION}"
+RUN echo "CUDA_VARIANT: ${CUDA_VARIANT}"
RUN echo "LD_LIBRARY_PATH: ${LD_LIBRARY_PATH}"
diff --git a/cd/python/docker/python_images.sh b/cd/python/docker/python_images.sh
index e82f642..4d9cf13 100755
--- a/cd/python/docker/python_images.sh
+++ b/cd/python/docker/python_images.sh
@@ -46,13 +46,13 @@
build() {
# use this flag to set LD_LIBRARY_PATH for cuda files
if [[ $mxnet_variant = cu* ]]; then
- cuda_version=${mxnet_variant:2:-1}.${mxnet_variant: -1}
+ cuda_variant=${mxnet_variant:2:-1}.${mxnet_variant: -1}
else
- cuda_version="false"
+ cuda_variant="false"
fi
# NOTE: Ensure the correct context root is passed in when building - Dockerfile expects ./wheel_build
- docker build -t "${image_name}" --build-arg BASE_IMAGE="${base_image}" --build-arg MXNET_COMMIT_ID=${GIT_COMMIT} --build-arg MXNET_VARIANT=${mxnet_variant} --build-arg CUDA_VERSION=${cuda_version} -f ${resources_path}/Dockerfile ./wheel_build
+ docker build -t "${image_name}" --build-arg BASE_IMAGE="${base_image}" --build-arg MXNET_COMMIT_ID=${GIT_COMMIT} --build-arg MXNET_VARIANT=${mxnet_variant} --build-arg CUDA_VARIANT=${cuda_variant} -f ${resources_path}/Dockerfile ./wheel_build
}
test() {