remove normal head scale from cosine sim attention

Author: lucidrains

memory_efficient_attention_pytorch/memory_efficient_cosine_sim_attention.py

@@ -27,9 +27,6 @@ def attention(
     attn_bias = None,
     **kwargs
 ):
-    scale = q.shape[-1] ** -0.5
-    q = q * scale
-
     sim = einsum('b h i d, b h j d -> b h i j', q, k)
 
     if exists(attn_bias):
@@ -84,9 +81,6 @@ def numerically_unstable_memory_efficient_attention(
     k_bucket_size = 1024,
     eps = 1e-8
 ):
-    scale = q.shape[-1] ** -0.5
-    q = q * scale
-
     # chunk all the inputs
 
     q_chunks = q.split(q_bucket_size, dim = -2)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'memory-efficient-attention-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.8',
+  version = '0.0.9',
   license='MIT',
   description = 'Memory Efficient Attention - Pytorch',
   author = 'Phil Wang',