From 28d6a4399796a2dac707bffc5abf828269244405 Mon Sep 17 00:00:00 2001 From: DEKHTIARJonathan Date: Thu, 28 Apr 2022 14:18:29 -0700 Subject: [PATCH] [TF-TRT] Adding `tearDown()` to unittest base function. --- .../tensorrt/test/tf_trt_integration_test_base.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tensorflow/python/compiler/tensorrt/test/tf_trt_integration_test_base.py b/tensorflow/python/compiler/tensorrt/test/tf_trt_integration_test_base.py index dc1ab08fd27517..327072432b9f99 100644 --- a/tensorflow/python/compiler/tensorrt/test/tf_trt_integration_test_base.py +++ b/tensorflow/python/compiler/tensorrt/test/tf_trt_integration_test_base.py @@ -155,6 +155,15 @@ def setUp(self): if not is_tensorrt_enabled(): self.skipTest("Test requires TensorRT") + def tearDown(self): + """Making sure to clean artifact.""" + idx = 0 + while gc.garbage: + gc.collect() # Force GC to destroy the TRT engine cache. + idx += 1 + if idx >= 10: # After 10 iterations, break to avoid infinite collect. + break + def _GetTensorSpec(self, shape, mask, dtype, name): # Set dimension i to None if mask[i] == False assert len(shape) == len(mask)