From 2fbfa0ab4ed5ff25bc46aa7f589159280a9a68c4 Mon Sep 17 00:00:00 2001
From: Adil Zouitine <adilzouitinegm@gmail.com>
Date: Tue, 18 May 2021 12:48:32 +0200
Subject: [PATCH] Replaces the composition of fft1D into one fft2D

---
 fnet.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fnet.py b/fnet.py
index 04eae94..2473816 100644
--- a/fnet.py
+++ b/fnet.py
@@ -27,7 +27,7 @@ def __init__(self):
     super().__init__()
 
   def forward(self, x):
-    x = torch.fft.fft(torch.fft.fft(x, dim=-1), dim=-2).real
+    x = torch.fft.fft2(x).real
     return x
 
 class FNet(nn.Module):