Fix TF Rag OOM issue (#24122)

fix Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

Fix TF Rag OOM issue (#24122)
fix Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
707023d1 · Yih-Dar · GitHub · f2b91835 · 707023d1
Unverified Commit 707023d1 authored Jun 09, 2023 by Yih-Dar Committed by GitHub Jun 09, 2023
Show whitespace changes
Inline Side-by-side

Showing with 6 additions and 0 deletions

tests/models/rag/test_modeling_tf_rag.py tests/models/rag/test_modeling_tf_rag.py +6 -0

No files found.
--- a/tests/models/rag/test_modeling_tf_rag.py
+++ b/tests/models/rag/test_modeling_tf_rag.py
 from __future__ import annotations

+import gc
 import json
 import os
 import shutil
@@ -550,6 +551,11 @@ class TFRagDPRBartTest(TFRagTestMixin, unittest.TestCase):
 @require_sentencepiece
 @require_tokenizers
 class TFRagModelIntegrationTests(unittest.TestCase):
+    def tearDown(self):
+        super().tearDown()
+        # clean-up as much as possible GPU memory occupied by PyTorch
+        gc.collect()
+
    @cached_property
    def token_model(self):
        return TFRagTokenForGeneration.from_pretrained_question_encoder_generator(