Encoding with model in float16 leads to "mat1 and mat2 must have the same dtype" error

[ahmedkooli]

Issue:

Hey! When loading the model dunzhang/stella_en_400M_v5 in torch.float16 and encoding a text, I run into this error: "RuntimeError: mat1 and mat2 must have the same dtype, but got Half and Float".

Code to reproduce:

from sentence_transformers import SentenceTransformer
import torch

model_name = "dunzhang/stella_en_400M_v5"
dev = torch.device("cuda")
model = SentenceTransformer(model_name, device=dev, trust_remote_code=True, 
                            model_kwargs={"torch_dtype": torch.float16})#.half()

sentences = ["This is a sentence."]
output = model.encode(sentences)

Uncommenting the .half() will fix the problem. However this problem doesn't appear with other models such as sentence-transformers/all-MiniLM-L6-v2.

Versions:

sentence-transformers==3.0.1
torch==2.4.0+cu121

Thanks for your help

[ir2718]

Hi,

this looks like a bug to me. What happens is if the loaded model contains any class other than Transformer, the parameters in it won't be cast to the torch_dtype value. Have a look (here).

[ahmedkooli]

Hey, thanks for your answer. What would you suggest? Should the code be changed to

module = module_class.load(module_path, **kwargs)

[ir2718]

@ahmedkooli

Have a look at the PR i added, that should help.