fix causal masking in sparse conv attention

lucidrains · Feb 5, 2021 · 40f4119 · 40f4119
1 parent 503139a
commit 40f4119
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/dalle_pytorch/attention.py b/dalle_pytorch/attention.py
@@ -143,7 +143,8 @@ def forward(self, x, mask = None):
 
         # mask image attention
 
-        mask = rearrange(img_seq, 'i -> () i ()') <= k_img_indices
+        q_img_indices = rearrange(img_seq, 'i -> () i ()')
+        mask =  q_img_indices >= k_img_indices
 
         # image can attend to all of text
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'dalle-pytorch',
   packages = find_packages(),
-  version = '0.0.52',
+  version = '0.0.53',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',