Fixed minor bug

when unloading vae during TI training, generating images after
training will error out

Author: Fampai

modules/textual_inversion/textual_inversion.py

@@ -409,6 +409,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
 
     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