diff options
| author | Void Agent <void@jayrup.hermes> | 2026-08-02 13:29:24 +0100 |
|---|---|---|
| committer | Void Agent <void@jayrup.hermes> | 2026-08-02 13:29:24 +0100 |
| commit | 799df091c7a163607cca1d5dcafb664f915ca847 (patch) | |
| tree | de4193167b6c5f3902d2c940beae95ca5d8b6283 | |
| parent | 8c2dc42642b1cdbfb5dfacd9c188821b03c84c4b (diff) | |
loss_reweight: CPU generators for randint/randperm (CUDA generators unsupported on torch 2.4); cuda-path verified
| -rw-r--r-- | src/loss_reweight.py | 9 |
1 files changed, 5 insertions, 4 deletions
diff --git a/src/loss_reweight.py b/src/loss_reweight.py index 3274807..bcfdf07 100644 --- a/src/loss_reweight.py +++ b/src/loss_reweight.py @@ -53,7 +53,7 @@ def _weighted_loss(logits, y, mode, q_id, batch_k, V, device): if mode == 'q': w[y.view(-1) == q_id] = WEIGHT elif mode == 'ctrl_random': - g = torch.Generator(device=device).manual_seed(1000 + batch_k) + g = torch.Generator().manual_seed(1000 + batch_k) # CPU generator (randperm) n_q = int((y == q_id).sum().item()) flat = torch.arange(y.numel(), device=device) non_q = flat[y.view(-1) != q_id] @@ -88,9 +88,10 @@ def train(mode, seed, max_iters, batch_size): bs = batch_size blk = args['block_size'] V = args['vocab_size'] - # identical minibatch order for every model: per-seed generator, fixed start - g = torch.Generator(device=device).manual_seed(20260731 + seed) - gval = torch.Generator(device=device).manual_seed(777 + seed) + # identical minibatch order for every model: per-seed CPU generator, fixed start + # (torch.randint does not accept CUDA generators on torch 2.4) + g = torch.Generator().manual_seed(20260731 + seed) + gval = torch.Generator().manual_seed(777 + seed) def get_batch(split): d = train_data if split == 'train' else val_data |