Update env.yml

env.yml

@@ -4,6 +4,7 @@ channels:
   - defaults
 dependencies:
   - python
+  - pip
   - pip:
     - datasets
     - huggingface-hub
@@ -33,5 +34,16 @@ dependencies:
     - dill
     - pox
     - pathos
-    - "jax[cuda_local] jaxlib[cuda_local] --extra-index-url https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
-    - "torch torchaudio torchvision triton xformers --extra-index-url https://download.pytorch.org/whl/cu121"
+  - pip:
+    - '--index-url=https://storage.googleapis.com/jax-releases/jax_cuda_releases.html'
+    - jax[cuda_local]
+  - pip:
+    - '--index-url=https://storage.googleapis.com/jax-releases/jax_cuda_releases.html'
+    - jaxlib[cuda_local]
+  - pip:
+    - '--index-url=https://download.pytorch.org/whl/cu121'
+    - torch
+    - torchaudio
+    - torchvision
+    - triton
+    - xformers