fix: precision during training

Author: stephantul

model2vec/train/base.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from typing import Any, TypeVar
 
 import numpy as np
@@ -11,6 +12,8 @@
 
 from model2vec import StaticModel
 
+logger = logging.getLogger(__name__)
+
 
 class FinetunableStaticModel(nn.Module):
     def __init__(self, *, vectors: torch.Tensor, tokenizer: Tokenizer, out_dim: int = 2, pad_id: int = 0) -> None:
@@ -26,9 +29,16 @@ def __init__(self, *, vectors: torch.Tensor, tokenizer: Tokenizer, out_dim: int
         self.pad_id = pad_id
         self.out_dim = out_dim
         self.embed_dim = vectors.shape[1]
-        self.vectors = vectors
 
-        self.embeddings = nn.Embedding.from_pretrained(vectors.clone().float(), freeze=False, padding_idx=pad_id)
+        self.vectors = vectors
+        if self.vectors.dtype != torch.float32:
+            dtype = str(self.vectors.dtype)
+            logger.warning(
+                f"Your vectors are {dtype} precision, converting to to torch.float32 to avoid compatibility issues."
+            )
+            self.vectors = vectors.float()
+
+        self.embeddings = nn.Embedding.from_pretrained(vectors.clone(), freeze=False, padding_idx=pad_id)
         self.head = self.construct_head()
         self.w = self.construct_weights()
         self.tokenizer = tokenizer