Merge pull request #88 from TheMattBin/master

Update ocr.py
VikParuchuri · May 8, 2024 · ef5a333 · ef5a333
2 parents 0bb1de4 + 2b25cdd
commit ef5a333
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/surya/ocr.py b/surya/ocr.py
@@ -62,7 +62,7 @@ def run_recognition(images: List[Image.Image], langs: List[List[str]], rec_model
 
 def run_ocr(images: List[Image.Image], langs: List[List[str]], det_model, det_processor, rec_model, rec_processor, batch_size=None) -> List[OCRResult]:
     det_predictions = batch_text_detection(images, det_model, det_processor)
-    if det_model.device == "cuda":
+    if det_model.device.type == "cuda":
         torch.cuda.empty_cache() # Empty cache from first model run
 
     slice_map = []