Error in model, scaling only q matrix not qK.T dot product (qk.T/sqrt(dim_per_head)

As per Vaswani et al, 2017 p.4
Is torch.matmul(q, k.transpose(2, 3)) / math.sqrt(dim_per_head) not q / math.sqrt(dim_per_head)
https://arxiv.org/pdf/1912.05372.pdf

Author: BenoitDalFerro

xlm/model/transformer.py

@@ -208,8 +208,8 @@ def unshape(x):
                     k, v = cache[self.layer_id]
             cache[self.layer_id] = (k, v)
 
-        q = q / math.sqrt(dim_per_head)                                       # (bs, n_heads, qlen, dim_per_head)
         scores = torch.matmul(q, k.transpose(2, 3))                           # (bs, n_heads, qlen, klen)
+        scores = scores / math.sqrt(dim_per_head)                             # (bs, n_heads, qlen, klen)
         mask = (mask == 0).view(mask_reshape).expand_as(scores)               # (bs, n_heads, qlen, klen)
         scores.masked_fill_(mask, -float('inf'))                              # (bs, n_heads, qlen, klen)