add missing device and dtype to SD1UNet's UpBlocks

This commit is contained in:
Laurent 2024-09-08 13:47:35 +00:00 committed by Laureηt
parent 444882a734
commit 602c63ce76

View file

@ -188,7 +188,7 @@ class SD1UNet(fl.Chain):
fl.UseContext(context="unet", key="residuals").compose(lambda x: x[-1]),
MiddleBlock(device=device, dtype=dtype),
),
UpBlocks(),
UpBlocks(device=device, dtype=dtype),
fl.Chain(
fl.GroupNorm(channels=320, num_groups=32, device=device, dtype=dtype),
fl.SiLU(),