From d68017037ce0f42dc036b30e1de36d8c49b73429 Mon Sep 17 00:00:00 2001 From: Robert Sachunsky Date: Mon, 9 Dec 2024 11:27:11 +0000 Subject: [PATCH] do_prediction: trigger GC to avoid CUDA OOM --- src/eynollah/eynollah.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/eynollah/eynollah.py b/src/eynollah/eynollah.py index 50f0f34..90824c8 100644 --- a/src/eynollah/eynollah.py +++ b/src/eynollah/eynollah.py @@ -999,7 +999,7 @@ class Eynollah: prediction_true = prediction_true.astype(np.uint8) #del model - #gc.collect() + gc.collect() return prediction_true def do_padding_with_scale(self,img, scale): @@ -1209,6 +1209,7 @@ class Eynollah: img_patch = np.zeros((n_batch_inference, img_height_model, img_width_model, 3)) prediction_true = prediction_true.astype(np.uint8) + gc.collect() return prediction_true def extract_page(self):