Unverified Commit b894b780 authored by Xiaomeng Zhao's avatar Xiaomeng Zhao Committed by GitHub
Browse files

Merge pull request #1570 from myhloli/dev

fix(magic_pdf): limit batch ratio for GPU memory
parents 9da857d5 db8be974
......@@ -177,7 +177,11 @@ def doc_analyze(
if torch.cuda.is_available() and device != 'cpu' or npu_support:
gpu_memory = get_vram(device)
if gpu_memory is not None and gpu_memory >= 7:
batch_ratio = int((gpu_memory-3) // 1.5)
# batch_ratio = int((gpu_memory-3) // 1.5)
batch_ratio = 2
if 8 < gpu_memory:
batch_ratio = 4
if batch_ratio >= 1:
logger.info(f'gpu_memory: {gpu_memory} GB, batch_ratio: {batch_ratio}')
batch_model = BatchAnalyze(model=custom_model, batch_ratio=batch_ratio)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment