From 890e68aaf75ae80d5eb2fa95b4bf1adf78b96881 Mon Sep 17 00:00:00 2001
From: Fampai <>
Date: Mon, 31 Oct 2022 10:07:12 -0400
Subject: [PATCH] Fixed minor bug

when unloading vae during TI training, generating images after
training will error out
---
 modules/textual_inversion/textual_inversion.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index 54a734f13..0aeb0459b 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -409,6 +409,7 @@ Last saved image: {html.escape(last_saved_image)}<br/>
 
     filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt')
     save_embedding(embedding, checkpoint, embedding_name, filename, remove_cached_checksum=True)
+    shared.sd_model.first_stage_model.to(devices.device)
 
     return embedding, filename