From 68ae8d856a3cb2a606ee0ed212c113606d0b35a0 Mon Sep 17 00:00:00 2001
From: pingsutw <pingsutw@gmail.com>
Date: Thu, 27 Feb 2020 10:05:06 +0800
Subject: [PATCH] Fix gpu not found when running TVM docker

---
 docker/build.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/build.sh b/docker/build.sh
index a2e21d140f6f..21141b71ca7d 100755
--- a/docker/build.sh
+++ b/docker/build.sh
@@ -79,7 +79,7 @@ if [ "$#" -lt 1 ] || [ ! -e "${SCRIPT_DIR}/Dockerfile.${CONTAINER_TYPE}" ]; then
 fi
 
 # Use nvidia-docker if the container is GPU.
-if [[ "${DOCKER_IMAGE_NAME}" == *"gpu"* ]]; then
+if [[ "${CONTAINER_TYPE}" == *"gpu"* ]]; then
     if ! type "nvidia-docker" 1> /dev/null 2> /dev/null
     then
         DOCKER_BINARY="docker"