step_emb sized to K_MAX+2 (aligned evals unroll up to 64)

ptr/models.py

@@ -75,8 +75,9 @@ class PtrTransformer(nn.Module):
         self.state_emb = nn.Embedding(S, d_model)
         self.pos_emb = nn.Embedding(max_len, d_model)
         self.k_emb = nn.Embedding(self.K_MAX + 2, d_model)
-        # step embeddings up to 24 unrolls (2x the max training unroll)
-        self.step_emb = nn.Embedding(24, d_model)
+        # step embeddings must cover the largest unroll count used anywhere:
+        # aligned models evaluate at R = k, up to k = K_MAX (64)
+        self.step_emb = nn.Embedding(self.K_MAX + 2, d_model)
         self.blocks = nn.ModuleList(
             Block(d_model, n_heads, d_ff, zero_init=zero_init)
             for _ in range(depth)