🐛 Fix MDSI runtime on CUDA

Mean_cuda is not implemented for complex types (pytorch/pytorch#46982)
francois-rozet · Dec 10, 2020 · 0e89000 · 0e89000
1 parent d277888
commit 0e89000
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/piqa/__init__.py b/piqa/__init__.py
@@ -5,4 +5,4 @@
 specific image quality assessement metric.
 """
 
-__version__ = '1.0.5'
+__version__ = '1.0.6'
diff --git a/piqa/mdsi.py b/piqa/mdsi.py
@@ -112,7 +112,9 @@ def mdsi(
 
     # Mean deviation similarity
     gcs_q = gcs ** q
-    score = (gcs_q - gcs_q.mean((-1, -2), keepdim=True)).abs()
+    gcs_q_avg = torch.view_as_real(gcs_q).mean((-2, -3), True)
+    gcs_q_avg = torch.view_as_complex(gcs_q_avg)
+    score = (gcs_q - gcs_q_avg).abs()
     mds = (score ** rho).mean((-1, -2)) ** (o / rho)
 
     return mds