Free memory after model export (#800)

huggingface · Jul 4, 2024 · e7ffee9 · e7ffee9
1 parent 50d867c
commit e7ffee9
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py
@@ -12,6 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import gc
 import logging
 import warnings
 from pathlib import Path
@@ -28,6 +29,8 @@
 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:
     from optimum.intel.openvino.configuration import OVConfig
@@ -367,6 +370,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:
         torch.cuda.is_available = orig_cuda_check