Update install.sh

MolCrafts · Dec 2, 2024 · e28f88c · e28f88c
1 parent 2982dda
commit e28f88c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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