protect against row being all masked out in sdpa

lucidrains · Nov 2, 2024 · d46e9cc · d46e9cc
1 parent 899b414
commit d46e9cc
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/nGPT_pytorch/nGPT.py b/nGPT_pytorch/nGPT.py
@@ -264,6 +264,8 @@ def forward(
         # for non-autoregressive masking
 
         if exists(mask):
+            row_all_masked_out = ~mask.any(dim = -1)
+
             mask = rearrange(mask, 'b j -> b 1 1 j')
 
             if exists(self.mask_value):
@@ -282,6 +284,9 @@ def forward(
         out = self.merge_heads(out)
         out = self.to_out(out)
 
+        if exists(mask) and row_all_masked_out.any():
+            out = einx.where('b n, b n d, -> b n d', ~row_all_masked_out, out, 0.)
+
         if not return_values:
             return out
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "nGPT-pytorch"
-version = "0.2.4"
+version = "0.2.5"
 description = "nGPT"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }