From 66b8b77feb5fafd3696044ee33a80f26cb693a13 Mon Sep 17 00:00:00 2001 From: CaptainJack2491 Date: Mon, 17 Aug 2026 16:15:28 +0100 Subject: feat(gpu): add CUDA support, chunked eval, bisect prime search, and scaling to 100k (48 tests green) --- src/model_api.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) (limited to 'src/model_api.py') diff --git a/src/model_api.py b/src/model_api.py index 31ffe31..85b5d55 100644 --- a/src/model_api.py +++ b/src/model_api.py @@ -32,10 +32,16 @@ def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int Inputs are LEFT-padded to the global layout so positions match training exactly (batch/singleton invariance — codex BLOCKER fix).""" + try: + device = next(model.parameters()).device + except StopIteration: + device = x.device + if x.device != device: + x = x.to(device) x = pad_inputs(x, cfg) max_len = max_len or cfg.max_out_len B = x.shape[0] - y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=x.device) + y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device) for _ in range(max_len): out = model(x, y_in) nxt = out["logits"][:, -1].argmax(-1) -- cgit v1.2.3