diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py index 927c98ac37..757244df55 100644 --- 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