summaryrefslogtreecommitdiffstats
path: root/training/strategy/ti.py
diff options
context:
space:
mode:
authorVolpeon <git@volpeon.ink>2023-03-23 11:07:57 +0100
committerVolpeon <git@volpeon.ink>2023-03-23 11:07:57 +0100
commit0767c7bc82645186159965c2a6be4278e33c6721 (patch)
treea136470ab85dbb99ab51d9be4a7831fe21612ab3 /training/strategy/ti.py
parentFix (diff)
downloadtextual-inversion-diff-0767c7bc82645186159965c2a6be4278e33c6721.tar.gz
textual-inversion-diff-0767c7bc82645186159965c2a6be4278e33c6721.tar.bz2
textual-inversion-diff-0767c7bc82645186159965c2a6be4278e33c6721.zip
Update
Diffstat (limited to 'training/strategy/ti.py')
-rw-r--r--training/strategy/ti.py4
1 files changed, 2 insertions, 2 deletions
diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index bd0d178..2038e34 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -156,8 +156,8 @@ def textual_inversion_strategy_callbacks(
156 @torch.no_grad() 156 @torch.no_grad()
157 def on_sample(step): 157 def on_sample(step):
158 with ema_context(): 158 with ema_context():
159 unet_ = accelerator.unwrap_model(unet, False) 159 unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True)
160 text_encoder_ = accelerator.unwrap_model(text_encoder, False) 160 text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True)
161 161
162 orig_unet_dtype = unet_.dtype 162 orig_unet_dtype = unet_.dtype
163 orig_text_encoder_dtype = text_encoder_.dtype 163 orig_text_encoder_dtype = text_encoder_.dtype