X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=beaver.py;h=5ee468e0bf070a79ce1deff3045ab77b43e1becd;hb=9fb6b9a0f1aff25ec1b559c318102060178f3a90;hp=4f41832fcd928b7f25bc9722a66b2d810229711e;hpb=eb86c22964f03f186ee225f129bc260128b10f9a;p=beaver.git

diff --git a/beaver.py b/beaver.py
index 4f41832..5ee468e 100755
--- a/beaver.py
+++ b/beaver.py
@@ -524,7 +524,10 @@ amm_generator = None
 if args.noncausal_prompt:
     amm_generator = lambda d: torch.logical_and(
         torch.arange(d)[None, None, :, None] < torch.arange(d)[None, None, None, :],
-        torch.arange(d)[None, None, :, None] >= d // 2,
+        torch.logical_or(
+            torch.arange(d)[None, None, :, None] >= d // 2,
+            torch.arange(d)[None, None, None, :] >= d // 2,
+        ),
     )
 
 model = mygpt.MyGPT(