Update CUDA version.

The older CUDA version for some reason fails to compile the extensions on current Pytorch version.
SINGROUP · Apr 23, 2024 · 6cce1ce · 6cce1ce
1 parent 1807794
commit 6cce1ce
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/environment.yml b/environment.yml
@@ -1,15 +1,14 @@
 name: ml-spm
 channels:
   - pytorch
-  - nvidia/label/cuda-11.8.0
+  - nvidia/label/cuda-12.1.0
   - conda-forge
   - defaults
 dependencies:
   - python=3.11
   - pip
   - pytorch
-  - torchvision
-  - torchaudio
+  - pytorch-cuda=12.1
   - cuda
   - ninja
   - matplotlib