stability measure number 2

Author: lucidrains

dalle_pytorch/attention.py

@@ -45,7 +45,9 @@ def forward(self, x, mask = None):
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
 
-        dots = torch.einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
+        q = q * self.scale
+
+        dots = torch.einsum('b h i d, b h j d -> b h i j', q, k)
         mask_value = max_neg_value(dots)
 
         if exists(mask):

setup.py

@@ -4,7 +4,7 @@
   name = 'dalle-pytorch',
   packages = find_packages(),
   include_package_data = True,
-  version = '0.12.1',
+  version = '0.12.2',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',