From 4800c62f6ed598accb950a8b715a6cab8a264e1e Mon Sep 17 00:00:00 2001
From: Sofie Van Landeghem <svlandeg@users.noreply.github.com>
Date: Wed, 18 Feb 2026 01:03:46 +0100
Subject: [PATCH] Fix MockModel's device definition (#535)

* fix MockModel's device definition

* cleanup
---
 tests/test_engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_engine.py b/tests/test_engine.py
index 01591111..784ffcb9 100644
--- a/tests/test_engine.py
+++ b/tests/test_engine.py
@@ -31,7 +31,7 @@ class MockModel:
     def __init__(self, vocab_size=262):  # 256 bytes + 6 special tokens
         self.vocab_size = vocab_size
         self.config = MockConfig()
-        self._device = "cpu"
+        self._device = torch.device("cpu")
 
     def get_device(self):
         return self._device