address lucidrains#292

ZijianDu · Dec 23, 2023 · 5578ac4 · 5578ac4
1 parent d446a41
commit 5578ac4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/setup.py b/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',

diff --git a/vit_pytorch/simple_vit_with_fft.py b/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)