From 8f093fb799246f7dd9104ff44728da0c53a9f67a Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Sat, 10 Jun 2023 09:47:38 +0200 Subject: [PATCH] Avoid OOM in doctest CI (#24139) * fix * fix --------- Co-authored-by: ydshieh --- src/transformers/testing_utils.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py index 0a35494765..f703575190 100644 --- a/src/transformers/testing_utils.py +++ b/src/transformers/testing_utils.py @@ -1870,9 +1870,18 @@ def preprocess_string(string, skip_cuda_tests): ): is_cuda_found = True break + modified_string = "" if not is_cuda_found: modified_string = "".join(codeblocks) + + if ">>>" in modified_string: + lines = modified_string.split("\n") + indent = len(lines[-1]) - len(lines[-1].lstrip()) + + cleanup = ">>> import gc; gc.collect() # doctest: +IGNORE_RESULT" + modified_string += "\n" + " " * indent + cleanup + return modified_string