Merge pull request #140 from Isalia20/proper_gpu_memory_release

Proper gpu memory release
nadermx · Apr 19, 2024 · 22596ef · 22596ef
2 parents c590858 + dfa4dd8
commit 22596ef
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/backgroundremover/u2net/detect.py b/backgroundremover/u2net/detect.py
@@ -151,5 +151,6 @@ def predict(net, item):
         img = Image.fromarray(predict_np * 255).convert("RGB")
 
         del d1, d2, d3, d4, d5, d6, d7, pred, predict, predict_np, inputs_test, sample
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
 
         return img