move releasing cache outside of the loop

2026-04-08 22:35:09 +00:00 · 2023-07-30 13:30:42 +09:00
parent 54a4aa22ed
commit e6034b7eb6
1 changed files with 3 additions and 2 deletions
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -2165,8 +2165,9 @@ def cache_batch_latents(
            info.latents = latent
            if flip_aug:
                info.latents_flipped = flipped_latent
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+                
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()


 def cache_batch_text_encoder_outputs(