Use gpu version of pytorch, add nvidia channel, add pytorch-cuda, tor…

…chvision
radionets-project · Aug 5, 2024 · fd75e56 · fd75e56
1 parent 3e7e0f4
commit fd75e56
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/environment.yml b/environment.yml
@@ -5,14 +5,17 @@ channels:
     - pytorch
     - defaults
     - conda-forge
+    - nvidia
 dependencies:
   - python
-  - pytorch
+  - pytorch-gpu
+  - pytorch-cuda=12.1
+  - torchvision
   - cudatoolkit
   - cartopy
   - numpy
   - numba
+  - towncrier
   - pip
   - pip:
-    - towncrier
     - -e .