From 44d84bd8b145d8e8c4313b06b49e4c744ba2b130 Mon Sep 17 00:00:00 2001
From: Jacky <18255193+kthui@users.noreply.github.com>
Date: Tue, 7 May 2024 12:25:44 -0700
Subject: [PATCH] Update the Conda installation command to match the latest
 from PyTorch (#129)

---
 tools/gen_pb_exec_env.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tools/gen_pb_exec_env.sh b/tools/gen_pb_exec_env.sh
index f89f0e2..d7fe86b 100755
--- a/tools/gen_pb_exec_env.sh
+++ b/tools/gen_pb_exec_env.sh
@@ -41,7 +41,7 @@ export PYTHONNOUSERSITE=True
 conda install -c conda-forge libstdcxx-ng=12 -y
 
 # install PyTorch
-conda install pytorch torchvision torchaudio pytorch-cuda -c pytorch -c nvidia -y
+conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia -y
 
 # pack environment
 rm -f pb_exec_env_model.py.tar.gz