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:
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user