From 9fb6b9a0f1aff25ec1b559c318102060178f3a90 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Fri, 24 Mar 2023 20:51:50 +0100 Subject: [PATCH] Update --- beaver.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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( -- 2.39.5