norm not needed when reusing attention in lookvit

Gzzgz · Jul 19, 2024 · ec6c48b · ec6c48b
1 parent 547bf94
commit ec6c48b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -6,7 +6,7 @@
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '1.7.1',
+  version = '1.7.2',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   long_description=long_description,

diff --git a/vit_pytorch/look_vit.py b/vit_pytorch/look_vit.py
@@ -77,7 +77,7 @@ def __init__(
 
         self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
 
-        self.norm = LayerNorm(dim)
+        self.norm = LayerNorm(dim) if not reuse_attention else nn.Identity()
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)