added pytorch-rocm dependency

Helsinki-NLP · Jul 5, 2024 · bec8aed · bec8aed
1 parent 57e0036
commit bec8aed
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/envs/hf.yml b/envs/hf.yml
@@ -11,4 +11,6 @@ dependencies:
     - transformers==4.38.2
     - sentencepiece==0.1.99
     - accelerate==0.31.0
-    - torch==2.2.1
+    - torch==2.2.1+rocm5.3
+    - torchvision==0.15.2+rocm5.3
+    - torchaudio==0.13.0+rocm5.3