From 6bc1bd2ec42c4787d5a12c12bb294519021a736c Mon Sep 17 00:00:00 2001 From: Nikos Livathinos Date: Tue, 10 Dec 2024 15:05:00 +0000 Subject: [PATCH] fix: Correct the way to set GPU for EasyOCR, RapidOCR Signed-off-by: Nikos Livathinos Signed-off-by: Christoph Auer --- docling/models/easyocr_model.py | 2 +- docling/models/rapid_ocr_model.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docling/models/easyocr_model.py b/docling/models/easyocr_model.py index 68beb2fa..a10f5ba2 100644 --- a/docling/models/easyocr_model.py +++ b/docling/models/easyocr_model.py @@ -48,7 +48,7 @@ class EasyOcrModel(BaseOcrModel): use_gpu = any( filter( lambda x: str(x) in device, - [AcceleratorDevice.CUDA, AcceleratorDevice.MPS], + ["cuda", "mps"], ) ) diff --git a/docling/models/rapid_ocr_model.py b/docling/models/rapid_ocr_model.py index 7cf8c7d3..ec471253 100644 --- a/docling/models/rapid_ocr_model.py +++ b/docling/models/rapid_ocr_model.py @@ -42,7 +42,7 @@ class RapidOcrModel(BaseOcrModel): # Decide the accelerator devices device = decide_device(accelerator_options.device) - use_cuda = AcceleratorDevice.CUDA in device + use_cuda = "cuda" in device use_dml = accelerator_options.device == AcceleratorDevice.AUTO intra_op_num_threads = accelerator_options.num_threads