fix: use pytorch+cuda version. test on Ubuntu & Windows

Lambdaoda · Jun 20, 2024 · 8464f72 · 8464f72
1 parent 31c91b1
commit 8464f72
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -1,3 +1,5 @@
+--find-links https://download.pytorch.org/whl/torch_stable.html
+
 accelerate==0.28.0
 audio-separator==0.17.2
 av==12.1.0
@@ -20,8 +22,8 @@ opencv-python-headless==4.9.0.80
 opencv-python==4.9.0.80
 pillow==10.3.0
 setuptools==70.0.0
-torch==2.2.2
-torchvision==0.17.2
+torch==2.2.2+cu121
+torchvision==0.17.2+cu121
 tqdm==4.66.4
 transformers==4.39.2
 xformers==0.0.25.post1