Aquila-Consortium/map2map
32
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix again rfft call broken by pytorch breaking BC

Browse Source
This commit is contained in:
Yin Li 2021-04-22 13:32:13 -04:00
parent 795cefe38e
commit 5e4633b125
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
map2map/models/power.py
View File

@ -19,10 +19,11 @@ def power(x):
try: try:
x = torch.fft.rfftn(x, s=signal_size) # new version broke BC x = torch.fft.rfftn(x, s=signal_size) # new version broke BC
P = x.real.square() + x.imag.square()
except AttributeError: except AttributeError:
x = torch.rfft(x, signal_ndim) 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.mean(dim=0)
P = P.sum(dim=0) P = P.sum(dim=0)
del x del x