From 002e24992500b4b16b79ae0225b336bf949c611b Mon Sep 17 00:00:00 2001
From: Yin Li <eelregit@gmail.com>
Date: Mon, 12 Apr 2021 12:43:29 -0400
Subject: [PATCH] Fix rfft call broken by pytorch breaking BC

---
 map2map/models/power.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/map2map/models/power.py b/map2map/models/power.py
index d65acc1..9b72d2a 100644
--- a/map2map/models/power.py
+++ b/map2map/models/power.py
@@ -13,10 +13,15 @@ def power(x):
     frequency of the input.
     """
     signal_ndim = x.dim() - 2
-    kmax = min(d for d in x.shape[-signal_ndim:]) // 2
+    signal_size = x.shape[-signal_ndim:]
+    kmax = min(s for s in signal_size) // 2
     even = x.shape[-1] % 2 == 0
 
-    x = torch.rfft(x, signal_ndim)
+    try:
+        x = torch.fft.rfftn(x, s=signal_size)  # new version broke BC
+    except AttributeError:
+        x = torch.rfft(x, signal_ndim)
+
     P = x.pow(2).sum(dim=-1)
     P = P.mean(dim=0)
     P = P.sum(dim=0)