fix mask and heads issue

lucidrains · Aug 23, 2021 · 2f1fb7c · 2f1fb7c
1 parent 15e72e5
commit 2f1fb7c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/invariant_point_attention/invariant_point_attention.py b/invariant_point_attention/invariant_point_attention.py
@@ -145,6 +145,7 @@ def forward(
 
         if exists(mask):
             mask = rearrange(mask, 'b i -> b i ()') * rearrange(mask, 'b j -> b () j')
+            mask = repeat(mask, 'b i j -> (b h) i j', h = h)
             mask_value = max_neg_value(attn_logits)
             attn_logits = attn_logits.masked_fill(~mask, mask_value)
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'invariant-point-attention',
   packages = find_packages(),
-  version = '0.1.2',
+  version = '0.1.4',
   license='MIT',
   description = 'Invariant Point Attention',
   author = 'Phil Wang',