diff --git a/environment.yml b/environment.yml
index 4845ed0..7fdbead 100644
--- a/environment.yml
+++ b/environment.yml
@@ -26,9 +26,6 @@ dependencies:
   - line_profiler
   - memory_profiler
   - snakeviz
-  - pytorch
-  - torchvision
-  - torchaudio
   - plotly
   - nbmake
   - h5py
@@ -39,3 +36,6 @@ dependencies:
     - build
     - tensorboard
     - pytorch-lightning
+    - torch
+    - torchvision
+    - torchaudio