Fix CI GPU install problem

rapidsai · Oct 9, 2024 · efbc828 · efbc828
1 parent bd20039
commit efbc828
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/ci/test_gpu.sh b/ci/test_gpu.sh
@@ -21,15 +21,12 @@ mamba install \
   conda-forge::pytorch>=2.0=*cuda* \
   conda-forge::transformers \
   conda-forge::pytest \
+  conda-forge::sentence-transformers  \
+  conda-forge::sentencepiece \
   -c conda-forge \
   --override-channels \
   --yes
 
-# Have to install sentence-transformers from pip
-# because conda-forge leads to a torch vision conflict
-# which leads to it being installed on CPUs
-pip3 install sentence-transformers sentencepiece
-
 # Install the crossfit package in editable mode with test dependencies
 pip3 install -e '.[test]'
 # Running tests