From 1a4221c2ba62bec311ae27ca4940c6fe38b1cef9 Mon Sep 17 00:00:00 2001
From: Brian Jou <brian@withpersona.com>
Date: Mon, 19 Aug 2024 09:20:05 -0700
Subject: [PATCH] fix: resolves memory leak caused by using CRAFT detector with
 detect() or readtext()

---
 easyocr/detection.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/easyocr/detection.py b/easyocr/detection.py
index ff96fef6b1d..4a553d23861 100644
--- a/easyocr/detection.py
+++ b/easyocr/detection.py
@@ -58,6 +58,9 @@ def test_net(
     with torch.no_grad():
         y, feature = net(x)
 
+    # remove x and feature from device, whether GPU or CPU
+    del x, feature
+
     boxes_list, polys_list = [], []
     for out in y:
         # make score and link map
@@ -82,6 +85,11 @@ def test_net(
         boxes_list.append(boxes)
         polys_list.append(polys)
 
+    # remove y from device, whether GPU or CPU, and check if cuda was used before calling empty_cache() to clean up
+    del y
+    if device == 'cuda':
+        torch.cuda.empty_cache()
+
     return boxes_list, polys_list
 
 def get_detector(trained_model, device='cpu', quantize=True, cudnn_benchmark=False):