From 5e4633b125863955718bff0cbed990210945b9bb Mon Sep 17 00:00:00 2001 From: Yin Li Date: Thu, 22 Apr 2021 13:32:13 -0400 Subject: [PATCH] Fix again rfft call broken by pytorch breaking BC --- map2map/models/power.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/map2map/models/power.py b/map2map/models/power.py index 9b72d2a..81fb7ac 100644 --- a/map2map/models/power.py +++ b/map2map/models/power.py @@ -19,10 +19,11 @@ def power(x): try: x = torch.fft.rfftn(x, s=signal_size) # new version broke BC + P = x.real.square() + x.imag.square() except AttributeError: x = torch.rfft(x, signal_ndim) + P = x.square().sum(dim=-1) - P = x.pow(2).sum(dim=-1) P = P.mean(dim=0) P = P.sum(dim=0) del x