Parcourir la source

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

myhloli il y a 3 semaines
Parent
commit
39271b45de
1 fichiers modifiés avec 1 ajouts et 0 suppressions
  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 = []