do_prediction: trigger GC to avoid CUDA OOM

2 weeks ago · d68017037c
parent ad748d0039
commit d68017037c
1 changed files with 2 additions and 1 deletions
--- 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):