From e28f88c56848fa2f981dad55541b495a36321c4c Mon Sep 17 00:00:00 2001
From: Jichen Li <42854324+Roy-Kid@users.noreply.github.com>
Date: Mon, 2 Dec 2024 15:10:14 +0100
Subject: [PATCH] Update install.sh

---
 .devcontainer/features/molpot/install.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.devcontainer/features/molpot/install.sh b/.devcontainer/features/molpot/install.sh
index fd42716..42390ea 100644
--- a/.devcontainer/features/molpot/install.sh
+++ b/.devcontainer/features/molpot/install.sh
@@ -5,8 +5,8 @@ echo "Activating feature 'pytorch'"
 
 which pip > /dev/null || (apt update && apt install python3-pip -y -qq)
 
-if [ -z "${MOLPOTENABLECUDA}" ]; then
-    echo "Installing pytorch with ${MOLPOTENABLECUDA} support"
+if [ "${PYTORCHCUDAVERSION}" ]; then
+    echo "Installing pytorch with ${PYTORCHCUDAVERSION} support"
     pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${PYTORCHCUDAVERSION}
     echo "Pytorch with CUDA support installed!"
 else