mem masks should be on the continuous transformer wrapper

lucidrains · Jan 16, 2024 · aa380f1 · aa380f1
1 parent f185da8
commit aa380f1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.27.6',
+  version = '1.27.7',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

diff --git a/x_transformers/continuous.py b/x_transformers/continuous.py
@@ -84,6 +84,7 @@ def forward(
         mask = None,
         return_attn = False,
         mems = None,
+        mem_masks = None,
         pos = None,
         prepend_embeds = None,
         prepend_mask = None,
@@ -125,7 +126,7 @@ def forward(
 
         # attention layers
 
-        x, intermediates = self.attn_layers(x, mask = mask, mems = mems, return_hiddens = True, **kwargs)
+        x, intermediates = self.attn_layers(x, mask = mask, mems = mems, mem_masks = mem_masks, return_hiddens = True, **kwargs)
 
         # splice out memory tokens