fix mixed precision training on train_dreambooth_inpaint_lora (#3138)

cast to weight dtype

fix mixed precision training on train_dreambooth_inpaint_lora (#3138)
cast to weight dtype
0ddc5bf7 · Lucca Zenóbio · GitHub · c5933c9c · 0ddc5bf7
Unverified Commit 0ddc5bf7 authored Apr 25, 2023 by Lucca Zenóbio Committed by GitHub Apr 25, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

examples/research_projects/dreambooth_inpaint/train_dreambooth_inpaint_lora.py ...jects/dreambooth_inpaint/train_dreambooth_inpaint_lora.py +1 -1

No files found.
--- a/examples/research_projects/dreambooth_inpaint/train_dreambooth_inpaint_lora.py
+++ b/examples/research_projects/dreambooth_inpaint/train_dreambooth_inpaint_lora.py
@@ -735,7 +735,7 @@ def main():
                        torch.nn.functional.interpolate(mask, size=(args.resolution // 8, args.resolution // 8))
                        for mask in masks
                    ]
-                )
+                ).to(dtype=weight_dtype)
                mask = mask.reshape(-1, 1, args.resolution // 8, args.resolution // 8)

                # Sample noise that we'll add to the latents