1 files changed, 15 insertions, 0 deletions
diff --git a/tests/test_loss_reweight.py b/tests/test_loss_reweight.py
index 7df66c9..aa1b69c 100644
--- a/tests/test_loss_reweight.py
+++ b/tests/test_loss_reweight.py
@@ -47,7 +47,22 @@ def test_weighting_semantics():
     assert abs(l_r1.item() - l_r2.item()) < 1e-9, "ctrl_random must be deterministic"
 
 
+def test_forward_full_logits():
+    """Regression: model(X, Y)[0] must be full (B, T, V) logits.
+    Karpathy nanoGPT returns LAST-position logits only when targets=None."""
+    from model import GPT, GPTConfig
+    torch.manual_seed(0)
+    cfg = GPTConfig(n_layer=2, n_head=2, n_embd=16, block_size=32, bias=False,
+                    vocab_size=65, dropout=0.0)
+    m = GPT(cfg).to(DEV)
+    x = torch.randint(0, 65, (2, 32), device=DEV)
+    y = torch.randint(0, 65, (2, 32), device=DEV)
+    logits = m(x, y)[0]
+    assert logits.shape == (2, 32, 65), f"full logits expected, got {logits.shape}"
+
+
 if __name__ == '__main__':
     test_batch_determinism()
     test_weighting_semantics()
+    test_forward_full_logits()
     print(f"PASS ({DEV})")