norm not needed when reusing attention in lookvit

setup.py

@@ -6,7 +6,7 @@ with open('README.md') as f:
 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,

vit_pytorch/look_vit.py

@@ -77,7 +77,7 @@ class Attention(Module):
 
         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)