| 245 | |
| 246 | class ConvResblock(nn.Module): |
| 247 | def __init__(self, in_features=320, out_features=320) -> None: |
| 248 | super().__init__() |
| 249 | self.f_t = nn.Linear(1280, out_features * 2) |
| 250 | |
| 251 | self.gn_1 = nn.GroupNorm(32, in_features) |
| 252 | self.f_1 = nn.Conv2d(in_features, out_features, kernel_size=3, padding=1) |
| 253 | |
| 254 | self.gn_2 = nn.GroupNorm(32, out_features) |
| 255 | self.f_2 = nn.Conv2d(out_features, out_features, kernel_size=3, padding=1) |
| 256 | |
| 257 | skip_conv = in_features != out_features |
| 258 | self.f_s = nn.Conv2d(in_features, out_features, kernel_size=1, padding=0) if skip_conv else nn.Identity() |
| 259 | |
| 260 | def forward(self, x, t): |
| 261 | x_skip = x |