free memory after model export

huggingface · Jul 4, 2024 · c32e065 · c32e065
1 parent 384dda5
commit c32e065
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py
@@ -14,6 +14,7 @@
 
 import logging
 import warnings
+import gc
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 
@@ -27,6 +28,7 @@
 from optimum.exporters.openvino.convert import export_from_model
 from optimum.intel.utils.import_utils import is_openvino_tokenizers_available, is_transformers_version
 from optimum.utils.save_utils import maybe_load_preprocessors
+from .utils import clear_class_registry
 
 
 if TYPE_CHECKING:
@@ -366,6 +368,10 @@ class StoreAttr(object):
 
     if convert_tokenizer:
         maybe_convert_tokenizers(library_name, output, model, preprocessors)
+
+    clear_class_registry()
+    del model
+    gc.collect()
 
     # Unpatch modules after GPTQ export
     if do_gptq_patching: