add cuda empty cache to properly release cuda memory

nadermx · Apr 19, 2024 · dfa4dd8 · dfa4dd8
1 parent 4ba6485
commit dfa4dd8
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