浏览代码

fix(magic_pdf): correct batch ratio conditions for GPU memory

- Update conditions for batch ratio assignment:
  -8 <= gpu_memory < 10: batch_ratio = 2 - 10 <= gpu_memory <= 12: batch_ratio =4
- This fix ensures proper batch ratio selection for GPU memory sizes
myhloli 10 月之前
父节点
当前提交
b6710b9988
共有 1 个文件被更改,包括 2 次插入2 次删除
  1. 2 2
      magic_pdf/model/doc_analyze_by_custom_model.py

+ 2 - 2
magic_pdf/model/doc_analyze_by_custom_model.py

@@ -178,9 +178,9 @@ def doc_analyze(
         gpu_memory = int(os.getenv("VIRTUAL_VRAM_SIZE", round(get_vram(device))))
         if gpu_memory is not None and gpu_memory >= 8:
 
-            if 8 <= gpu_memory <= 10:
+            if 8 <= gpu_memory < 10:
                 batch_ratio = 2
-            elif 10 < gpu_memory <= 12:
+            elif 10 <= gpu_memory <= 12:
                 batch_ratio = 4
             elif 12 < gpu_memory <= 16:
                 batch_ratio = 8