Update ocr.py

VikParuchuri · May 5, 2024 · 2b25cdd · 2b25cdd
1 parent ac07e7c
commit 2b25cdd
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) -> 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 = []