loss_reweight: pass Y to forward for full logits (Karpathy nanoGPT returns last-position logits when targets=None); regression test

2 files changed, 18 insertions, 3 deletions

diff --git a/src/loss_reweight.py b/src/loss_reweight.py
index bcfdf07..1bf3124 100644
--- a/src/loss_reweight.py
+++ b/src/loss_reweight.py
@@ -108,8 +108,8 @@ def train(mode, seed, max_iters, batch_size):
             for _ in range(50):
                 X, Y = get_batch('val')
                 with torch.no_grad():
-                    logits = model(X)[0]
-                lv.append(F.cross_entropy(logits.view(-1, V), Y.view(-1)).item())
+                    _, loss = model(X, Y)
+                lv.append(loss.item())
             v = np.mean(lv)
             model.train()
             if v < best_val:
@@ -119,7 +119,7 @@ def train(mode, seed, max_iters, batch_size):
             if it % 1000 == 0:
                 print(f"  iter {it}: val={v:.4f}")
         X, Y = get_batch('train')
-        logits = model(X)[0]
+        logits = model(X, Y)[0]  # pass Y: targets=None would give last-position logits only
         loss = _weighted_loss(logits, Y, mode, q_id, it, V, device)
         loss.backward()
         opt.step()
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})")