diff options
| -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 |