From 8a073ee49f6ec29429c66b8ccfabba49bf167ebc Mon Sep 17 00:00:00 2001
From: Pam <pamhome21@gmail.com>
Date: Mon, 17 Jul 2023 17:51:26 +0500
Subject: [PATCH] vram leak fix

---
 library/train_util.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/library/train_util.py b/library/train_util.py
index a4021c4f..ea63b71d 100644
--- 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(