Fix MockModel's device definition (#535)

* fix MockModel's device definition * cleanup
2026-02-18 01:03:46 +01:00
parent f5fe7925ed
commit cac43e8511
1 changed files with 1 additions and 1 deletions
@@ -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