mirror of
https://github.com/finegrain-ai/refiners.git
synced 2024-11-09 23:12:02 +00:00
Init dtype and device correctly for OutputBlock
This commit is contained in:
parent
3a10baa9f8
commit
44e184d4d5
|
@ -242,7 +242,7 @@ class OutputBlock(fl.Chain):
|
||||||
|
|
||||||
def __init__(self, device: Device | str | None = None, dtype: DType | None = None) -> None:
|
def __init__(self, device: Device | str | None = None, dtype: DType | None = None) -> None:
|
||||||
super().__init__(
|
super().__init__(
|
||||||
fl.GroupNorm(channels=320, num_groups=32),
|
fl.GroupNorm(channels=320, num_groups=32, device=device, dtype=dtype),
|
||||||
fl.SiLU(),
|
fl.SiLU(),
|
||||||
fl.Conv2d(in_channels=320, out_channels=4, kernel_size=3, stride=1, padding=1, device=device, dtype=dtype),
|
fl.Conv2d(in_channels=320, out_channels=4, kernel_size=3, stride=1, padding=1, device=device, dtype=dtype),
|
||||||
)
|
)
|
||||||
|
|
Loading…
Reference in a new issue