diff --git a/docker/base/build_image.sh b/docker/base/build_image.sh
index 101dfebad..03a3185c4 100755
--- a/docker/base/build_image.sh
+++ b/docker/base/build_image.sh
@@ -4,7 +4,7 @@ SCRIPT_LOCATION=$0
 cd "$(dirname "$SCRIPT_LOCATION")"
 WORK_DIR=$(pwd)
 
-BASE_IMAGE="nvidia/cuda:11.8.0-devel-ubuntu22.04"
+BASE_IMAGE="nvidia/cuda:11.8.0-runtime-ubuntu22.04"
 IMAGE_NAME="eosphorosai/dbgpt"
 # zh: https://pypi.tuna.tsinghua.edu.cn/simple
 PIP_INDEX_URL="https://pypi.org/simple"
@@ -14,7 +14,7 @@ BUILD_LOCAL_CODE="false"
 LOAD_EXAMPLES="true"
 
 usage () {
-    echo "USAGE: $0 [--base-image nvidia/cuda:11.8.0-devel-ubuntu22.04] [--image-name db-gpt]"
+    echo "USAGE: $0 [--base-image nvidia/cuda:11.8.0-runtime-ubuntu22.04] [--image-name db-gpt]"
     echo "  [-b|--base-image base image name] Base image name"
     echo "  [-n|--image-name image name] Current image name, default: db-gpt"
     echo "  [-i|--pip-index-url pip index url] Pip index url, default: https://pypi.org/simple"
diff --git a/setup.py b/setup.py
index b12c35dfa..a9ee31213 100644
--- a/setup.py
+++ b/setup.py
@@ -216,8 +216,8 @@ def torch_requires(
         cuda_version = get_cuda_version()
         if not cuda_version:
             torch_pkgs = [
-                f"torch=={torch_version}+cpu",
-                f"torchvision=={torchvision_version}+cpu",
+                f"torch=={torch_version}",
+                f"torchvision=={torchvision_version}",
                 f"torchaudio=={torchaudio_version}",
             ]
         else: