diff --git a/docling/models/easyocr_model.py b/docling/models/easyocr_model.py index 7ad79b96..68beb2fa 100644 --- a/docling/models/easyocr_model.py +++ b/docling/models/easyocr_model.py @@ -45,10 +45,13 @@ class EasyOcrModel(BaseOcrModel): if self.options.use_gpu: device = decide_device(accelerator_options.device) # Enable easyocr GPU if running on CUDA, MPS - use_gpu = device in [ - AcceleratorDevice.CUDA, - AcceleratorDevice.MPS, - ] + use_gpu = any( + filter( + lambda x: str(x) in device, + [AcceleratorDevice.CUDA, AcceleratorDevice.MPS], + ) + ) + self.reader = easyocr.Reader( lang_list=self.options.lang, gpu=use_gpu, diff --git a/docling/models/rapid_ocr_model.py b/docling/models/rapid_ocr_model.py index f476b067..7cf8c7d3 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 = device == AcceleratorDevice.CUDA + use_cuda = AcceleratorDevice.CUDA in device use_dml = accelerator_options.device == AcceleratorDevice.AUTO intra_op_num_threads = accelerator_options.num_threads