mirror of
https://github.com/finegrain-ai/refiners.git
synced 2024-11-09 15:02:01 +00:00
add missing device and dtype to SD1UNet's UpBlocks
This commit is contained in:
parent
444882a734
commit
602c63ce76
|
@ -188,7 +188,7 @@ class SD1UNet(fl.Chain):
|
||||||
fl.UseContext(context="unet", key="residuals").compose(lambda x: x[-1]),
|
fl.UseContext(context="unet", key="residuals").compose(lambda x: x[-1]),
|
||||||
MiddleBlock(device=device, dtype=dtype),
|
MiddleBlock(device=device, dtype=dtype),
|
||||||
),
|
),
|
||||||
UpBlocks(),
|
UpBlocks(device=device, dtype=dtype),
|
||||||
fl.Chain(
|
fl.Chain(
|
||||||
fl.GroupNorm(channels=320, num_groups=32, device=device, dtype=dtype),
|
fl.GroupNorm(channels=320, num_groups=32, device=device, dtype=dtype),
|
||||||
fl.SiLU(),
|
fl.SiLU(),
|
||||||
|
|
Loading…
Reference in a new issue