Commit 053b0582 authored by Wael Ramadan's avatar Wael Ramadan
Browse files

feat: [#535] flush cuda memory

parent 2b775a8f
Pipeline #7969 passed with stage
in 19 minutes and 58 seconds
import gc
import logging
import os
import pathlib
......@@ -310,6 +311,9 @@ class MLP:
if "sentences" in analyzers and spans == "sentence":
document.fact_spans_to_sent()
gc.collect()
torch.cuda.empty_cache()
return document.to_json()
else:
return None
......@@ -393,4 +397,7 @@ class MLP:
result = doc.document_to_json(use_default_doc_path=False)
docs[actual_index] = result
gc.collect()
torch.cuda.empty_cache()
return docs
Supports Markdown
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