Fix DFT_series_decomp: use the top_k parameter

Author: rozsasarpi

models/TimeMixer.py

@@ -11,15 +11,15 @@ class DFT_series_decomp(nn.Module):
     Series decomposition block
     """
 
-    def __init__(self, top_k=5):
+    def __init__(self, top_k: int = 5):
         super(DFT_series_decomp, self).__init__()
         self.top_k = top_k
 
     def forward(self, x):
         xf = torch.fft.rfft(x)
         freq = abs(xf)
         freq[0] = 0
-        top_k_freq, top_list = torch.topk(freq, 5)
+        top_k_freq, top_list = torch.topk(freq, k=self.top_k)
         xf[freq <= top_k_freq.min()] = 0
         x_season = torch.fft.irfft(xf)
         x_trend = x - x_season