tests/rag/test_modeling_rag.py · 32adbb26d670c3226d1c874b2b59fc97be7e9445 · chenpangpang / transformers

Fix PyTorch RAG tests GPU OOM (#16881) · 32adbb26

Yih-Dar authored Apr 25, 2022



* add torch.cuda.empty_cache in some PT RAG tests

* torch.cuda.empty_cache in tearDownModule()

* tearDown()

* add gc.collect()
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

32adbb26

test_modeling_rag.py 44.4 KB

Replace test_modeling_rag.py