From 799df091c7a163607cca1d5dcafb664f915ca847 Mon Sep 17 00:00:00 2001 From: Void Agent Date: Sun, 2 Aug 2026 13:29:24 +0100 Subject: loss_reweight: CPU generators for randint/randperm (CUDA generators unsupported on torch 2.4); cuda-path verified --- src/loss_reweight.py | 9 +++++---- 1 file 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 -- cgit v1.2.3