fix multiheaded qk rmsnorm in nViT

Author: lucidrains

setup.py

@@ -6,7 +6,7 @@
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '1.8.4',
+  version = '1.8.5',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   long_description=long_description,

vit_pytorch/normalized_vit.py

@@ -76,8 +76,8 @@ def __init__(
 
         self.dropout = dropout
 
-        self.q_scale = nn.Parameter(torch.ones(dim_inner) * (dim_head ** 0.25))
-        self.k_scale = nn.Parameter(torch.ones(dim_inner) * (dim_head ** 0.25))
+        self.q_scale = nn.Parameter(torch.ones(heads, 1, dim_head) * (dim_head ** 0.25))
+        self.k_scale = nn.Parameter(torch.ones(heads, 1, dim_head) * (dim_head ** 0.25))
 
         self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
         self.merge_heads = Rearrange('b h n d -> b n (h d)')
@@ -90,15 +90,15 @@ def forward(
     ):
         q, k, v = self.to_q(x), self.to_k(x), self.to_v(x)
 
-        q = q * self.q_scale
-        k = k * self.k_scale
-
         q, k, v = map(self.split_heads, (q, k, v))
 
         # query key rmsnorm
 
         q, k = map(l2norm, (q, k))
 
+        q = q * self.q_scale
+        k = k * self.k_scale
+
         # scale is 1., as scaling factor is moved to s_qk (dk ^ 0.25) - eq. 16
 
         out = F.scaled_dot_product_attention(