From 4811a3d1df15f24b009e5ce168e1e5589e743d4b Mon Sep 17 00:00:00 2001 From: myhloli Date: Thu, 19 Sep 2024 20:42:11 +0800 Subject: [PATCH] fix(pdf-extract): ensure model is set to evaluation mode before processing Add model.eval() invocation to pdf_extract_kit initialization sequence to ensure the model is set to evaluation mode. This is critical for proper inference and performance metrics when processing PDF content. --- magic_pdf/model/pdf_extract_kit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index 406417cd..7f9c6808 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -64,6 +64,7 @@ def mfr_model_init(weight_dir, cfg_path, _device_='cpu'): task = tasks.setup_task(cfg) model = task.build_model(cfg) model = model.to(_device_) + model.eval() vis_processor = load_processor('formula_image_eval', cfg.config.datasets.formula_rec_eval.vis_processor.eval) mfr_transform = transforms.Compose([vis_processor, ]) return [model, mfr_transform]