Skip to content

Commit 4daa408

Browse files
authored
Pass on model_kwargs when loading a sentence-transformers model before export (#2126)
1 parent 3f00766 commit 4daa408

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

optimum/exporters/tasks.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2141,6 +2141,7 @@ def get_model_from_task(
21412141
use_auth_token = model_kwargs.pop("use_auth_token", None)
21422142
token = model_kwargs.pop("token", None)
21432143
trust_remote_code = model_kwargs.pop("trust_remote_code", False)
2144+
model_kwargs["torch_dtype"] = torch_dtype
21442145

21452146
if use_auth_token is not None:
21462147
warnings.warn(
@@ -2158,6 +2159,7 @@ def get_model_from_task(
21582159
token=token,
21592160
revision=revision,
21602161
trust_remote_code=trust_remote_code,
2162+
model_kwargs=model_kwargs,
21612163
)
21622164
else:
21632165
try:

0 commit comments

Comments
 (0)