(tvm) branch main updated: [Docker] Use Torch GPU on gpu device (#17676)

[yongwww]

This is an automated email from the ASF dual-hosted git repository.

yongwww pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new aeadc31c97 [Docker] Use Torch GPU on gpu device (#17676)
aeadc31c97 is described below

commit aeadc31c9714c4ba7878688d5ad7c0f19c1eb308
Author: Yong Wu <yongc...@gmail.com>
AuthorDate: Mon Feb 24 18:22:03 2025 -0800

    [Docker] Use Torch GPU on gpu device (#17676)
---
 docker/Dockerfile.ci_gpu              |  2 +-
 docker/install/ubuntu_install_onnx.sh | 45 +++++++++++++++++++++++++----------
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/docker/Dockerfile.ci_gpu b/docker/Dockerfile.ci_gpu
index 5fa0df9231..332cb9b4e0 100644
--- a/docker/Dockerfile.ci_gpu
+++ b/docker/Dockerfile.ci_gpu
@@ -89,7 +89,7 @@ COPY install/ubuntu_install_jax.sh 
/install/ubuntu_install_jax.sh
 RUN bash /install/ubuntu_install_jax.sh "cuda"
 
 COPY install/ubuntu_install_onnx.sh /install/ubuntu_install_onnx.sh
-RUN bash /install/ubuntu_install_onnx.sh
+RUN bash /install/ubuntu_install_onnx.sh "cuda"
 
 COPY install/ubuntu_install_libtorch.sh /install/ubuntu_install_libtorch.sh
 RUN bash /install/ubuntu_install_libtorch.sh
diff --git a/docker/install/ubuntu_install_onnx.sh 
b/docker/install/ubuntu_install_onnx.sh
index a8bebc2988..dc41c39d7c 100755
--- a/docker/install/ubuntu_install_onnx.sh
+++ b/docker/install/ubuntu_install_onnx.sh
@@ -30,6 +30,9 @@ set -o pipefail
 # Get the Python version
 PYTHON_VERSION=$(python3 -c "import sys; 
print(f'{sys.version_info.major}.{sys.version_info.minor}')")
 
+# Set default value for first argument
+DEVICE=${1:-cpu}
+
 # Install the onnx package
 pip3 install future
 
@@ -39,28 +42,46 @@ if [ "$PYTHON_VERSION" == "3.9" ]; then
         onnxruntime==1.19.2 \
         onnxoptimizer==0.2.7
 
-    pip3 install \
-        torch==2.6.0 \
-        torchvision==0.21.0 \
-        --extra-index-url https://download.pytorch.org/whl/cpu
+    if [ "$DEVICE" == "cuda" ]; then
+        pip3 install \
+            torch==2.6.0 \
+            torchvision==0.21.0
+    else
+        pip3 install \
+            torch==2.6.0 \
+            torchvision==0.21.0 \
+            --extra-index-url https://download.pytorch.org/whl/cpu
+    fi
 elif [ "$PYTHON_VERSION" == "3.11" ]; then
     pip3 install \
         onnx==1.17.0 \
         onnxruntime==1.20.1 \
         onnxoptimizer==0.2.7
 
-    pip3 install \
-        torch==2.6.0 \
-        torchvision==0.21.0 \
-        --extra-index-url https://download.pytorch.org/whl/cpu
+    if [ "$DEVICE" == "cuda" ]; then
+        pip3 install \
+            torch==2.6.0 \
+            torchvision==0.21.0
+    else
+        pip3 install \
+            torch==2.6.0 \
+            torchvision==0.21.0 \
+            --extra-index-url https://download.pytorch.org/whl/cpu
+    fi
 else
     pip3 install \
         onnx==1.12.0 \
         onnxruntime==1.12.1 \
         onnxoptimizer==0.2.7
 
-    pip3 install \
-        torch==2.4.1 \
-        torchvision==0.19.1 \
-        --extra-index-url https://download.pytorch.org/whl/cpu
+    if [ "$DEVICE" == "cuda" ]; then
+        pip3 install \
+            torch==2.4.1 \
+            torchvision==0.19.1
+    else
+        pip3 install \
+            torch==2.4.1 \
+            torchvision==0.19.1 \
+            --extra-index-url https://download.pytorch.org/whl/cpu
+    fi
 fi