Skip to content

Commit e7ffee9

Browse files
authored
Free memory after model export (#800)
1 parent 50d867c commit e7ffee9

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

optimum/exporters/openvino/__main__.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import gc
1516
import logging
1617
import warnings
1718
from pathlib import Path
@@ -28,6 +29,8 @@
2829
from optimum.intel.utils.import_utils import is_openvino_tokenizers_available, is_transformers_version
2930
from optimum.utils.save_utils import maybe_load_preprocessors
3031

32+
from .utils import clear_class_registry
33+
3134

3235
if TYPE_CHECKING:
3336
from optimum.intel.openvino.configuration import OVConfig
@@ -367,6 +370,10 @@ class StoreAttr(object):
367370
if convert_tokenizer:
368371
maybe_convert_tokenizers(library_name, output, model, preprocessors)
369372

373+
clear_class_registry()
374+
del model
375+
gc.collect()
376+
370377
# Unpatch modules after GPTQ export
371378
if do_gptq_patching:
372379
torch.cuda.is_available = orig_cuda_check

0 commit comments

Comments
 (0)