浏览代码

perf(model): adjust batch ratio for different GPU memory sizes

myhloli 9 月之前
父节点
当前提交
29e7a9488f
共有 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

@@ -182,9 +182,9 @@ def doc_analyze(
                 batch_ratio = 2
             elif 10 <= gpu_memory <= 12:
                 batch_ratio = 4
-            elif 12 < gpu_memory <= 16:
+            elif 12 < gpu_memory <= 20:
                 batch_ratio = 8
-            elif 16 < gpu_memory <= 32:
+            elif 20 < gpu_memory <= 32:
                 batch_ratio = 16
             else:
                 batch_ratio = 32