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

Add torch.inference mode
2024-10-28 12:03:07 -07:00
parent 7df99e6784 7703ae07a9
commit ba4e04fccb
1 changed files with 1 additions and 0 deletions
--- 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: