浏览代码

fix: adjust batch size calculation in prediction logic of predict_formula.py

myhloli 3 周之前
父节点
当前提交
39271b45de
共有 1 个文件被更改,包括 1 次插入0 次删除
  1. 1 0
      mineru/model/mfr/pp_formulanet_plus_m/predict_formula.py

+ 1 - 0
mineru/model/mfr/pp_formulanet_plus_m/predict_formula.py

@@ -131,6 +131,7 @@ class FormulaRecognizer(BaseOCRV20):
 
         if len(sorted_images) > 0:
             # 进行预测
+            batch_size = min(batch_size, max(1, 2 ** (len(sorted_images).bit_length() - 1))) if sorted_images else 1
             rec_formula = self.predict(sorted_images, batch_size)
         else:
             rec_formula = []