address #292

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '1.6.4',
+  version = '1.6.5',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   long_description_content_type = 'text/markdown',

vit_pytorch/simple_vit_with_fft.py

@@ -1,5 +1,5 @@
 import torch
-from torch.fft import fft
+from torch.fft import fft2
 from torch import nn
 
 from einops import rearrange, reduce, pack, unpack
@@ -128,7 +128,7 @@ def forward(self, img):
         device, dtype = img.device, img.dtype
 
         x = self.to_patch_embedding(img)
-        freqs = torch.view_as_real(fft(img))
+        freqs = torch.view_as_real(fft2(img))
 
         f = self.to_freq_embedding(freqs)