From d01c4ecefd3a6eaf3baab7fdcc2b0e2ea08e1f1b Mon Sep 17 00:00:00 2001
From: kthui <18255193+kthui@users.noreply.github.com>
Date: Mon, 6 May 2024 19:43:21 -0700
Subject: [PATCH] Update the Conda installation command to match the latest
 from PyTorch

---
 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