diff options
Diffstat (limited to 'tests/test_jlens_v3.py')
| -rw-r--r-- | tests/test_jlens_v3.py | 2 |
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): |