fix mixed precision training on train_dreambooth_inpaint_lora (#3138)

themrzmaster · web-flow · commit 0ddc5bf7b97e · 2023-04-25T15:22:57.000+05:30
cast to weight dtype
diff --git 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 collate_fn(examples):
                         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

Original file line number	Diff line number	Diff line change
`@@ -735,7 +735,7 @@ def collate_fn(examples):`
`735`	`735`	`torch.nn.functional.interpolate(mask, size=(args.resolution // 8, args.resolution // 8))`
`736`	`736`	`for mask in masks`
`737`	`737`	`]`
`738`		`- )`
	`738`	`+ ).to(dtype=weight_dtype)`
`739`	`739`	`mask = mask.reshape(-1, 1, args.resolution // 8, args.resolution // 8)`
`740`	`740`
`741`	`741`	`# Sample noise that we'll add to the latents`