implement CompVis#122 from compvis

Author: KAJdev

environment.yaml

@@ -7,14 +7,16 @@ dependencies:
   - pip=20.3
   - cudatoolkit=11.0
   - pytorch=1.7.0
-  - torchvision=0.12.0
-  - numpy=1.20.3
+  - torchvision=0.8.1
+  - numpy=1.19.2
   - pip:
     - albumentations==0.4.3
     - opencv-python==4.1.2.30
     - pudb==2019.2
     - imageio==2.9.0
     - imageio-ffmpeg==0.4.2
+    - kornia==0.6.6
+    - torchmetrics==v0.6.0
     - pytorch-lightning==1.4.2
     - omegaconf==2.1.1
     - test-tube>=0.7.5