vram leak fix

2026-04-10 06:54:17 +00:00 · 2023-07-17 17:51:26 +05:00
parent 7e20c6d1a1
commit 8a073ee49f
1 changed files with 2 additions and 0 deletions
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -2141,6 +2141,8 @@ def cache_batch_latents(
            info.latents = latent
            if flip_aug:
                info.latents_flipped = flipped_latent
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()


 def cache_batch_text_encoder_outputs(