From ae1a4b60c541e6befdb16767f0b04f10b5c8e5e1 Mon Sep 17 00:00:00 2001
From: HC Zhu <hczhu.mtv@gmail.com>
Date: Mon, 19 Jun 2023 21:13:14 -0700
Subject: [PATCH] update

---
 ai-ml/set_up_nvidia_gpu.sh | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)
 mode change 100644 => 100755 ai-ml/set_up_nvidia_gpu.sh

diff --git a/ai-ml/set_up_nvidia_gpu.sh b/ai-ml/set_up_nvidia_gpu.sh
old mode 100644
new mode 100755
index bd91f75..5a0f9f1
--- a/ai-ml/set_up_nvidia_gpu.sh
+++ b/ai-ml/set_up_nvidia_gpu.sh
@@ -7,13 +7,17 @@ installNvidiaDrive() {
   nvidia-smi
 }
 
+# See the latest image: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/index.html
+pytorch_image_version='23.02-py3'
+alias pytorch-docker="docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -it --rm nvcr.io/nvidia/pytorch:${pytorch_image_version}"
+
 installNvidiaLibContainer() {
   ubuntuVersion="$(lsb_release -a 2> /dev/null | grep 'Ubuntu .* LTS$' | awk ' { print $3 } ' | cut -d'.' -f1,2)"
   curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
   curl -s -L https://nvidia.github.io/libnvidia-container/ubuntu${ubuntuVersion}/libnvidia-container.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
   sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
   sudo nvidia-ctk runtime configure --runtime=docker
-  docker pull nvcr.io/nvidia/pytorch:23.05-py3
+  docker pull nvcr.io/nvidia/pytorch:${pytorch_image_version}
 }
 
 installNvidiaDrive
@@ -24,4 +28,4 @@ if [ -d "/usr/local/cuda-12" ]; then
   CUDA="$(ls /usr/local/cuda-12\.* -d | head -n1)"
   export PATH=${CUDA}/bin${PATH:+:${PATH}}
   export LD_LIBRARY_PATH=/usr/local/cuda-12.0/lib64:${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
-fi
+fi
\ No newline at end of file