From 15d84bcfb49a9ec5c281facf53bc023da5a4343e Mon Sep 17 00:00:00 2001
From: v0xie <28695009+v0xie@users.noreply.github.com>
Date: Fri, 28 Jun 2024 00:02:07 -0700
Subject: [PATCH] fix: add VRAM free after sample images

---
 train_network.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/train_network.py b/train_network.py
index 7ba073855..a5ef36462 100644
--- a/train_network.py
+++ b/train_network.py
@@ -1028,6 +1028,8 @@ def remove_model(old_ckpt_name):
 
                     self.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
 
+                    clean_memory_on_device(accelerator.device)
+
                     # 指定ステップごとにモデルを保存
                     if args.save_every_n_steps is not None and global_step % args.save_every_n_steps == 0:
                         accelerator.wait_for_everyone()
@@ -1084,6 +1086,8 @@ def remove_model(old_ckpt_name):
 
             self.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
 
+            clean_memory_on_device(accelerator.device)
+
             # end of epoch
 
         # metadata["ss_epoch"] = str(num_train_epochs)