Commit b4d4d0a0 authored by Raul Sirel's avatar Raul Sirel
Browse files

fallback to cpu on cuda oom error

parent de731cee
Pipeline #5879 canceled with stage
in 13 minutes and 4 seconds
......@@ -79,7 +79,6 @@ class MLP:
self.default_lang = default_language_code
self.use_default_lang = use_default_language_code
self.resource_dir = resource_dir
self.use_gpu = use_gpu
self.resource_dir_pathlib = pathlib.Path(resource_dir)
self.not_entities_path = self.resource_dir_pathlib / "concatenator" / "not_entities.txt"
......@@ -88,7 +87,13 @@ class MLP:
self.prepare_resources(refresh_data)
self.stanza_pipelines = self._load_stanza_pipelines(logging_level)
# This is for CUDA OOM exceptions. Fall back to CPU if needed.
try:
self.stanza_pipelines = self._load_stanza_pipelines(logging_level, use_gpu)
except RuntimeError:
# Try loading using CPU
self.stanza_pipelines = self._load_stanza_pipelines(logging_level, False)
self.entity_mapper = self._load_entity_mapper()
self.loaded_entity_files = []
......@@ -258,7 +263,7 @@ class MLP:
return [entity for sentence in pipeline.sentences for entity in sentence.entities]
def _load_stanza_pipelines(self, logging_level):
def _load_stanza_pipelines(self, logging_level, use_gpu):
"""
Initializes Stanza Pipeline objects all at once to save time later.
"""
......
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