amend

pytorch · Nov 2, 2023 · beee06c · beee06c
1 parent 7678d14
commit beee06c
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/.github/unittest/linux/scripts/install.sh b/.github/unittest/linux/scripts/install.sh
@@ -47,7 +47,10 @@ printf "* Installing tensordict\n"
 python setup.py develop
 
 # install torchsnapshot nightly
-python -m pip install git+https://github.com/pytorch/torchsnapshot --no-build-isolation
-
+if [[ "$TORCH_VERSION" == "nightly" ]]; then
+  python -m pip install git+https://github.com/pytorch/torchsnapshot --no-build-isolation
+elif [[ "$TORCH_VERSION" == "stable" ]]; then
+  python -m pip install torchsnapshot
+fi
 # smoke test
 python -c "import functorch;import torchsnapshot"