调整显存释放位置

2024-06-26 12:40:12 +08:00
parent 1240c4a884
commit 94d5c9b00f
1 changed files with 3 additions and 2 deletions
--- a/photo_review/photo_review.py
+++ b/photo_review/photo_review.py
@@ -172,6 +172,9 @@ def information_extraction(ie, phrecs):

            result = merge_result(result, ie_result)

+    # 完成一次任务释放显存
+    paddle.device.cuda.empty_cache()
+    logging.info("释放显存完成")
    return result


@@ -341,8 +344,6 @@ def main():
                session.execute(update_flag)
                session.commit()
                session.close()
-                # 完成一个案子释放显存
-                paddle.device.cuda.empty_cache()
        else:
            # 没有查询到新案子，等待一段时间后再查
            log = logging.getLogger()