fix mask for full attention so that last text token can actually predict the first image token

Author: lucidrains

dalle_pytorch/attention.py

@@ -48,7 +48,7 @@ def forward(self, x, mask = None):
         mask_value = max_neg_value(dots)
 
         if exists(mask):
-            mask = rearrange(mask, 'b i -> b () i ()') * rearrange(mask, 'b j -> b () () j')
+            mask = rearrange(mask, 'b j -> b () () j')
             dots.masked_fill_(~mask, mask_value)
             del mask
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'dalle-pytorch',
   packages = find_packages(),
-  version = '0.0.60',
+  version = '0.0.61',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',