dreambooth: fix #1566: maintain fp32 wrapper when saving a checkpoint to avoid crash when running fp16

Author: timh

examples/dreambooth/train_dreambooth.py

@@ -676,8 +676,8 @@ def collate_fn(examples):
                     if accelerator.is_main_process:
                         pipeline = DiffusionPipeline.from_pretrained(
                             args.pretrained_model_name_or_path,
-                            unet=accelerator.unwrap_model(unet),
-                            text_encoder=accelerator.unwrap_model(text_encoder),
+                            unet=accelerator.unwrap_model(unet, True),
+                            text_encoder=accelerator.unwrap_model(text_encoder, True),
                             revision=args.revision,
                         )
                         save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")