| 17 | |
| 18 | |
| 19 | class Block(nn.Module): |
| 20 | def __init__(self, n_in, n_out): |
| 21 | super().__init__() |
| 22 | self.conv = nn.Sequential(conv(n_in, n_out), nn.ReLU(), conv(n_out, n_out), nn.ReLU(), conv(n_out, n_out)) |
| 23 | self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() |
| 24 | self.fuse = nn.ReLU() |
| 25 | def forward(self, x): |
| 26 | return self.fuse(self.conv(x) + self.skip(x)) |
| 27 | |
| 28 | |
| 29 | def Encoder(latent_channels=4): |