fix mask needing repeats for non-autoregressive case

lucidrains · Aug 4, 2021 · cce1416 · cce1416
1 parent 575a416
commit cce1416
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/long_short_transformer/long_short_transformer.py b/long_short_transformer/long_short_transformer.py
@@ -167,7 +167,7 @@ def forward(self, x, mask = None):
             pkv = self.to_dynamic_proj(gkv)
 
             if exists(mask):
-                pmask = rearrange(mask, 'b (n s) -> b n s', s = s)
+                pmask = repeat(mask, 'b (n s) -> (b h) n s', s = s, h = h)
                 pkv.masked_fill_(~pmask[..., None], mask_value)
 
             pkv = pkv.softmax(dim = -2)

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'long-short-transformer',
   packages = find_packages(),
-  version = '0.0.4',
+  version = '0.0.5',
   license='MIT',
   description = 'Long Short Transformer - Pytorch',
   author = 'Phil Wang',