diff options
| -rw-r--r-- | src/loss_reweight.py | 6 | ||||
| -rw-r--r-- | tests/test_loss_reweight.py | 15 |
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})") |