From fc56dd007142264ef029eb8c0f9db11a6cf3d298 Mon Sep 17 00:00:00 2001 From: CaptainJack2491 Date: Sat, 29 Aug 2026 23:35:10 +0100 Subject: fix(transformer): increase positional embedding and causal mask buffer to 1024 to support worst-case probe traces --- src/models/transformer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/models/transformer.py b/src/models/transformer.py index a128ad9..4eb690a 100644 --- a/src/models/transformer.py +++ b/src/models/transformer.py @@ -32,13 +32,13 @@ class TransformerBaseline(PrimeModel): self.cfg = cfg d = cfg.d_model self.embed = nn.Embedding(cfg.vocab + 1, d) # +1 row = pad - self.pos = nn.Embedding(256, d) # supports sequences up to 256 + self.pos = nn.Embedding(1024, d) # supports full probe traces up to 1024 self.blocks = nn.ModuleList([CausalBlock(d, cfg.n_heads) for _ in range(cfg.n_layers)]) self.ln_f = nn.LayerNorm(d) self.head = nn.Linear(d, cfg.vocab) self.register_buffer( "causal_triu", - torch.triu(torch.ones(256, 256, dtype=torch.bool), diagonal=1), + torch.triu(torch.ones(1024, 1024, dtype=torch.bool), diagonal=1), persistent=False, ) -- cgit v1.2.3