fix(transformer): increase positional embedding and causal mask buffer to 1024 to support worst-case probe traces

1 files 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,
         )