Address nit

abheesht17 · Jul 26, 2023 · 74c0c27 · 74c0c27
1 parent 5cfb2c4
commit 74c0c27
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/keras_nlp/layers/modeling/f_net_encoder.py b/keras_nlp/layers/modeling/f_net_encoder.py
@@ -138,9 +138,9 @@ def call(self, inputs):
 
         def fourier_transform(input):
             # Apply FFT on the input and take the real part.
-            x = (input, ops.zeros_like(input))
-            mixing_output = ops.fft2(x)[0]
-            return mixing_output
+            real_in, imaginary_in = (input, ops.zeros_like(input))
+            real_out, _ = ops.fft2((real_in, imaginary_in))
+            return real_out
 
         def add_and_norm(input1, input2, norm_layer):
             return norm_layer(input1 + input2)