fix memory problem

huggingface · Mar 6, 2024 · b0818c0 · b0818c0
1 parent 8b3b521
commit b0818c0
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
@@ -79,11 +79,11 @@ def create_no_weights_model(self) -> None:
         # unlike Transformers, TXI won't accept any missing tensors so we need to materialize the model
         LOGGER.info(f"\t+ Loading no weights model from {self.no_weights_model}")
         with random_init_weights():
-            model = self.automodel_class.from_pretrained(
+            self.pretrained_model = self.automodel_class.from_pretrained(
                 self.no_weights_model, **self.config.hub_kwargs, device_map="auto", _fast_init=False
             )
         LOGGER.info("\t+ Saving no weights model")
-        model.save_pretrained(save_directory=self.no_weights_model)
+        self.pretrained_model.save_pretrained(save_directory=self.no_weights_model)
 
         if self.config.task in TEXT_GENERATION_TASKS:
             LOGGER.info("\t+ Modifying generation config for fixed length generation")