From 4daa40896f693649e21696c509cd98c7e0c40e3c Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 13 Dec 2024 17:03:26 +0100 Subject: [PATCH] Pass on `model_kwargs` when loading a sentence-transformers model before export (#2126) --- optimum/exporters/tasks.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py index 32e90c7da1..4db4130302 100644 --- 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: