From 8b67f20935e48b26c5803cf31e0e89b9cfaa22ab Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Mon, 29 Aug 2022 11:43:20 +0200 Subject: [PATCH] Fix memory leak issue in `torch_fx` tests (#18547) Co-authored-by: Lysandre Debut Co-authored-by: ydshieh --- tests/test_modeling_common.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 8f80d7fa42..5a2888a12a 100755 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -859,6 +859,10 @@ class ModelTesterMixin: f"serialized model {i}th output doesn't match model {i}th output for {model_class}", ) + # Avoid memory leak. Without this, each call increase RAM usage by ~20MB. + # (Even with this call, there are still memory leak by ~0.04MB) + self.clear_torch_jit_class_registry() + def test_headmasking(self): if not self.test_head_masking: return