From 0767c7bc82645186159965c2a6be4278e33c6721 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Thu, 23 Mar 2023 11:07:57 +0100 Subject: Update --- training/strategy/dreambooth.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) (limited to 'training/strategy/dreambooth.py') diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py index e5e84c8..28fccff 100644 --- a/training/strategy/dreambooth.py +++ b/training/strategy/dreambooth.py @@ -137,8 +137,8 @@ def dreambooth_strategy_callbacks( print("Saving model...") - unet_ = accelerator.unwrap_model(unet, False) - text_encoder_ = accelerator.unwrap_model(text_encoder, False) + unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=False) + text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=False) with ema_context(): pipeline = VlpnStableDiffusion( @@ -160,8 +160,8 @@ def dreambooth_strategy_callbacks( @torch.no_grad() def on_sample(step): with ema_context(): - unet_ = accelerator.unwrap_model(unet, False) - text_encoder_ = accelerator.unwrap_model(text_encoder, False) + unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True) + text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True) orig_unet_dtype = unet_.dtype orig_text_encoder_dtype = text_encoder_.dtype -- cgit v1.2.3-54-g00ecf