From cf6a73a7208ae17316a9287b5465702c8aee2fc9 Mon Sep 17 00:00:00 2001
From: hyh <2667004988@qq.com>
Date: Sun, 9 Jun 2024 13:24:48 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 itgpt-setup.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/itgpt-setup.sh b/itgpt-setup.sh
index bcdbc30..09be467 100755
--- a/itgpt-setup.sh
+++ b/itgpt-setup.sh
@@ -240,7 +240,7 @@ fi
 
 # Test / Install nvidia-docker
 if [[ ! -z "$NVIDIA_PRESENT" ]]; then
-    if sudo docker run --gpus all nvidia/cuda:11.0.3-base-ubuntu18.04 nvidia-smi &>/dev/null; then
+    if sudo docker run --rm --gpus all  nvidia/cuda:11.0.3-base-ubuntu18.04 nvidia-smi &>/dev/null; then
         echo "nvidia-docker 已启用并正常工作。退出脚本。"
     else
         echo "nvidia-docker 似乎没有启用。正在进行安装..."
@@ -249,7 +249,7 @@ if [[ ! -z "$NVIDIA_PRESENT" ]]; then
         curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
         sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
         sudo systemctl restart docker 
-        sudo docker run --gpus all nvidia/cuda:11.0.3-base-ubuntu18.04 nvidia-smi
+        sudo docker run --rm  --gpus all   nvidia/cuda:11.0.3-base-ubuntu18.04 nvidia-smi
         
     fi
 fi