Fix memory leak issue in torch_fx tests (#18547)

Co-authored-by: Lysandre Debut <hi@lysand.re>
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2022-08-29 11:43:20 +02:00
committed by GitHub
parent b10a3b3760
commit 8b67f20935

View File

@@ -859,6 +859,10 @@ class ModelTesterMixin:
f"serialized model {i}th output doesn't match model {i}th output for {model_class}", 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): def test_headmasking(self):
if not self.test_head_masking: if not self.test_head_masking:
return return