Update transformer to accelerator.device with it's weight_dtype

Viditnegi · web-flow · commit dd69d5434f3a · 2025-01-06T10:50:47.000+05:30
By default, the transformer always gets transferred to the device with float32 whether we specify the dtype as bfloat16 or float16. This fixes it.
diff --git a/examples/dreambooth/train_dreambooth_flux.py b/examples/dreambooth/train_dreambooth_flux.py
@@ -1157,6 +1157,7 @@ def main(args):
             "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
         )
 
+    transformer.to(accelerator.device, dtype=weight_dtype)
     vae.to(accelerator.device, dtype=weight_dtype)
     if not args.train_text_encoder:
         text_encoder_one.to(accelerator.device, dtype=weight_dtype)

Original file line number	Diff line number	Diff line change
`@@ -1157,6 +1157,7 @@ def main(args):`
`1157`	`1157`	`"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."`
`1158`	`1158`	`)`
`1159`	`1159`
	`1160`	`+ transformer.to(accelerator.device, dtype=weight_dtype)`
`1160`	`1161`	`vae.to(accelerator.device, dtype=weight_dtype)`
`1161`	`1162`	`if not args.train_text_encoder:`
`1162`	`1163`	`text_encoder_one.to(accelerator.device, dtype=weight_dtype)`