From 1fdfb2bc508249b68969aca5168ad14a013690c5 Mon Sep 17 00:00:00 2001 From: jane-temcious Date: Thu, 26 Jun 2025 06:49:07 +0530 Subject: [PATCH] Update base_pipeline.py Signed-off-by: jane-temcious --- docling/pipeline/base_pipeline.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/docling/pipeline/base_pipeline.py b/docling/pipeline/base_pipeline.py index 29475d68..4513df7c 100644 --- a/docling/pipeline/base_pipeline.py +++ b/docling/pipeline/base_pipeline.py @@ -93,8 +93,9 @@ class BasePipeline(ABC): pass def _unload(self, conv_res: ConversionResult): - pass - + # Clear GPU memory after processingAdd commentMore actions + clear_gpu_memory() + @classmethod @abstractmethod def get_default_options(cls) -> PipelineOptions: @@ -196,13 +197,18 @@ class PaginatedPipeline(BasePipeline): # TODO this is a bad name. return conv_res def _unload(self, conv_res: ConversionResult) -> ConversionResult: + # Unload page backends for page in conv_res.pages: if page._backend is not None: page._backend.unload() - + + # Unload input backend if conv_res.input._backend: conv_res.input._backend.unload() + # Call parent _unload to clean up GPU memory + super()._unload(conv_res) + return conv_res def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: