Finegrain/refiners
1
0
Fork 0
mirror of https://github.com/finegrain-ai/refiners.git synced 2024-09-19 19:05:28 +00:00
Code Issues Actions 9 Packages Projects Releases Wiki Activity

Init dtype and device correctly for OutputBlock

Browse source
This commit is contained in:
Doryan Kaced 2023-09-01 18:26:24 +02:00
parent 3a10baa9f8
commit 44e184d4d5
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
src/refiners/foundationals/latent_diffusion/stable_diffusion_xl/unet.py
View file

@ -242,7 +242,7 @@ class OutputBlock(fl.Chain):
def __init__(self, device: Device | str | None = None, dtype: DType | None = None) -> None:
super().__init__(
fl.GroupNorm(channels=320, num_groups=32),
fl.GroupNorm(channels=320, num_groups=32, device=device, dtype=dtype),
fl.SiLU(),
fl.Conv2d(in_channels=320, out_channels=4, kernel_size=3, stride=1, padding=1, device=device, dtype=dtype),
)