empry cache

huggingface · Mar 6, 2024 · 8fcd85f · 8fcd85f
1 parent b0818c0
commit 8fcd85f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/optimum_benchmark/backends/py_txi/backend.py b/optimum_benchmark/backends/py_txi/backend.py
@@ -84,6 +84,8 @@ def create_no_weights_model(self) -> None:
             )
         LOGGER.info("\t+ Saving no weights model")
         self.pretrained_model.save_pretrained(save_directory=self.no_weights_model)
+        del self.pretrained_model
+        torch.cuda.empty_cache()
 
         if self.config.task in TEXT_GENERATION_TASKS:
             LOGGER.info("\t+ Modifying generation config for fixed length generation")