fix an error thanks to @ShomyLiu at #2

flash_pytorch/flash_pytorch.py

@@ -339,8 +339,8 @@ def forward(
 
             lin_out = einsum('b g d e, b g n d -> b g n e', lin_kv, lin_q)
         else:
-            lin_kv = einsum('b g n d, b g n e -> b d e', lin_k, v) / n
-            lin_out = einsum('b g n d, b d e -> b g n e', lin_q, lin_kv)
+            lin_kv = einsum('b g n d, b g n e -> b g d e', lin_k, v) / n
+            lin_out = einsum('b g n d, b g d e -> b g n e', lin_q, lin_kv)
 
         # fold back groups into full sequence, and excise out padding
 

setup.py

@@ -3,11 +3,12 @@
 setup(
   name = 'FLASH-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.2',
+  version = '0.1.4',
   license='MIT',
   description = 'FLASH - Transformer Quality in Linear Time - Pytorch',
   author = 'Phil Wang',
   author_email = '[email protected]',
+  long_description_content_type = 'text/markdown',
   url = 'https://github.com/lucidrains/FLASH-pytorch',
   keywords = [
     'artificial intelligence',