Skip to content

GitLab

"...git@developer.sourcefind.cn:modelzoo/qwen_lmdeploy.git" did not exist on "af2f072e911a8797231ce92406d3e0c7e4d00163"
Unverified Commit 707023d1 authored by Yih-Dar's avatar Yih-Dar Committed by GitHub
Browse files

Fix TF Rag OOM issue (#24122) 



fix
Co-authored-by: default avatarydshieh <ydshieh@users.noreply.github.com>
parent f2b91835
Hide whitespace changes
Inline Side-by-side
Showing with 6 additions and 0 deletions
tests/models/rag/test_modeling_tf_rag.py
View file @ 707023d1
from __future__ import annotations from __future__ import annotations
import gc
import json import json
import os import os
import shutil import shutil
...@@ -550,6 +551,11 @@ class TFRagDPRBartTest(TFRagTestMixin, unittest.TestCase): ...@@ -550,6 +551,11 @@ class TFRagDPRBartTest(TFRagTestMixin, unittest.TestCase):
@require_sentencepiece @require_sentencepiece
@require_tokenizers @require_tokenizers
class TFRagModelIntegrationTests(unittest.TestCase): class TFRagModelIntegrationTests(unittest.TestCase):
def tearDown(self):
super().tearDown()
# clean-up as much as possible GPU memory occupied by PyTorch
gc.collect()
@cached_property @cached_property
def token_model(self): def token_model(self):
return TFRagTokenForGeneration.from_pretrained_question_encoder_generator( return TFRagTokenForGeneration.from_pretrained_question_encoder_generator(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment