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