Pass on model_kwargs when loading a sentence-transformers model bef…

…ore export (#2126)
huggingface · Dec 13, 2024 · 4daa408 · 4daa408
1 parent 3f00766
commit 4daa408
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -2141,6 +2141,7 @@ def get_model_from_task(
             use_auth_token = model_kwargs.pop("use_auth_token", None)
             token = model_kwargs.pop("token", None)
             trust_remote_code = model_kwargs.pop("trust_remote_code", False)
+            model_kwargs["torch_dtype"] = torch_dtype
 
             if use_auth_token is not None:
                 warnings.warn(
@@ -2158,6 +2159,7 @@ def get_model_from_task(
                 token=token,
                 revision=revision,
                 trust_remote_code=trust_remote_code,
+                model_kwargs=model_kwargs,
             )
         else:
             try: