| 234 | |
| 235 | |
| 236 | class ImageUnembedding(nn.Module): |
| 237 | def __init__(self, in_channels=320, out_channels=6) -> None: |
| 238 | super().__init__() |
| 239 | self.gn = nn.GroupNorm(32, in_channels) |
| 240 | self.f = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) |
| 241 | |
| 242 | def forward(self, x) -> torch.Tensor: |
| 243 | return self.f(F.silu(self.gn(x))) |
| 244 | |
| 245 | |
| 246 | class ConvResblock(nn.Module): |