test

huggingface · Dec 16, 2024 · 8c8d073 · 8c8d073
1 parent 31f96ba
commit 8c8d073
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optimum_benchmark/backends/py_txi/backend.py b/optimum_benchmark/backends/py_txi/backend.py
@@ -5,7 +5,7 @@
 import torch
 from huggingface_hub import snapshot_download
 from py_txi import TEI, TGI, TEIConfig, TGIConfig
-from safetensors.torch import save_file
+from safetensors.torch import save_file, save_model
 
 from ...task_utils import TEXT_EMBEDDING_TASKS, TEXT_GENERATION_TASKS
 from ..base import Backend
@@ -60,7 +60,7 @@ def create_no_weights_model(self) -> None:
             self.pretrained_model = self.automodel_loader.from_pretrained(
                 self.no_weights_model, **self.config.model_kwargs, device_map="auto", _fast_init=False
             )
-        save_file(tensors=self.pretrained_model.state_dict(), filename=filename, metadata={"format": "pt"})
+        self.pretrained_model.save_pretrained(save_directory=self.no_weights_model)
         del self.pretrained_model
         torch.cuda.empty_cache()