SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix TF Rag OOM issue (#24122)

Browse Source 
fix

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2023-06-09 15:03:11 +02:00
committed by GitHub
parent f2b918356c
commit 707023d155
1 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/models/rag/test_modeling_tf_rag.py
View File

@@ -1,5 +1,6 @@
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):
@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(