Skip to content

Commit

Permalink
[TF-TRT] Adding tearDown() to unittest base function.
Browse files Browse the repository at this point in the history
  • Loading branch information
DEKHTIARJonathan committed May 11, 2022
1 parent 03018aa commit 28d6a43
Showing 1 changed file with 9 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 28d6a43

Please sign in to comment.