From 98d46a3f085cc7529c9ddf3b7c4918cc248e9162 Mon Sep 17 00:00:00 2001 From: Anatoly Belikov Date: Mon, 18 Mar 2024 17:33:53 +0300 Subject: [PATCH] delete vae and text encoders after use in SDXL training script (#6693) delete vae and text encoders after use Co-authored-by: Sayak Paul --- examples/text_to_image/train_text_to_image_sdxl.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/text_to_image/train_text_to_image_sdxl.py b/examples/text_to_image/train_text_to_image_sdxl.py index 7028e82850..cb1feb806c 100644 --- a/examples/text_to_image/train_text_to_image_sdxl.py +++ b/examples/text_to_image/train_text_to_image_sdxl.py @@ -911,6 +911,7 @@ def main(args): ) precomputed_dataset = precomputed_dataset.with_transform(preprocess_train) + del compute_vae_encodings_fn, compute_embeddings_fn, text_encoder_one, text_encoder_two del text_encoders, tokenizers, vae gc.collect() torch.cuda.empty_cache()