Merge pull request #29 from aliencaocao/fix-inference-mode

yadong-lu · web-flow · commit ba4e04fccb67 · 2024-10-28T12:03:07.000-07:00
Add torch.inference mode
diff --git a/utils.py b/utils.py
@@ -64,6 +64,7 @@ def get_yolo_model(model_path):
     return model
 
 
+@torch.inference_mode()
 def get_parsed_content_icon(filtered_boxes, ocr_bbox, image_source, caption_model_processor, prompt=None):
     to_pil = ToPILImage()
     if ocr_bbox:
