瀏覽代碼

Merge pull request #636 from myhloli/dev

fix(pdf-extract): ensure model is set to evaluation mode before processing
Xiaomeng Zhao 1 年之前
父節點
當前提交
3a616cf29d
共有 1 個文件被更改,包括 1 次插入0 次删除
  1. 1 0
      magic_pdf/model/pdf_extract_kit.py

+ 1 - 0
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]