summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorVoid Agent <void@jayrup.hermes>2026-08-02 14:25:14 +0100
committerVoid Agent <void@jayrup.hermes>2026-08-02 14:25:14 +0100
commitd220be31e1e3082e56f9c59a735895839c99abce (patch)
treedd864686e71791f06afd5a81fa3c3b109f8753ff
parent31e1656f1fd4409c14108e63eb3e430f657746ac (diff)
test_jlens_v3: move tiny model to DEV (fix cpu/cuda mismatch on cuda runs)
-rw-r--r--tests/test_jlens_v3.py2
1 files changed, 1 insertions, 1 deletions
diff --git a/tests/test_jlens_v3.py b/tests/test_jlens_v3.py
index 5ac3f67..4ebeaf3 100644
--- a/tests/test_jlens_v3.py
+++ b/tests/test_jlens_v3.py
@@ -31,7 +31,7 @@ def _tiny_model(n_layer=2, d=16, block=32):
torch.manual_seed(0)
cfg = GPTConfig(n_layer=n_layer, n_head=2, n_embd=d, block_size=block,
bias=False, vocab_size=65, dropout=0.0)
- return GPT(cfg).eval()
+ return GPT(cfg).eval().to(DEV)
def _tiny_data(n=2000, block=32):