fix: move tensor to cpu (#86)

MinishLab · Oct 15, 2024 · 5c1ba32 · 5c1ba32
1 parent b4927e1
commit 5c1ba32
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/model2vec/model.py b/model2vec/model.py
@@ -111,7 +111,7 @@ def save_pretrained(self, path: PathLike, model_name: str | None = None) -> None
         """
         save_pretrained(
             folder_path=Path(path),
-            embeddings=self.embedding.weight.numpy(),
+            embeddings=self.embedding.weight.cpu().numpy(),
             tokenizer=self.tokenizer,
             config=self.config,
             base_model_name=self.base_model_name,