Skip to content

GitLab

Unverified Commit 158a5a87 authored by Muyang Li's avatar Muyang Li Committed by GitHub
Browse files

Remove the FP32 Wrapper when evaluating (#10617) 



Remove the FP32 Wrapper
Co-authored-by: default avatarLinoy Tsaban <57615435+linoytsaban@users.noreply.github.com>
parent 012d08b1
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
examples/dreambooth/train_dreambooth_flux.py
View file @ 158a5a87
...@@ -1716,9 +1716,9 @@ def main(args): ...@@ -1716,9 +1716,9 @@ def main(args):
pipeline = FluxPipeline.from_pretrained( pipeline = FluxPipeline.from_pretrained(
args.pretrained_model_name_or_path, args.pretrained_model_name_or_path,
vae=vae, vae=vae,
text_encoder=accelerator.unwrap_model(text_encoder_one), text_encoder=accelerator.unwrap_model(text_encoder_one, keep_fp32_wrapper=False),
text_encoder_2=accelerator.unwrap_model(text_encoder_two), text_encoder_2=accelerator.unwrap_model(text_encoder_two, keep_fp32_wrapper=False),
transformer=accelerator.unwrap_model(transformer), transformer=accelerator.unwrap_model(transformer, keep_fp32_wrapper=False),
revision=args.revision, revision=args.revision,
variant=args.variant, variant=args.variant,
torch_dtype=weight_dtype, torch_dtype=weight_dtype,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment