(self, in_channels=1024)
| 306 | # Also ConvResblock |
| 307 | class Upsample(nn.Module): |
| 308 | def __init__(self, in_channels=1024) -> None: |
| 309 | super().__init__() |
| 310 | self.f_t = nn.Linear(1280, in_channels * 2) |
| 311 | |
| 312 | self.gn_1 = nn.GroupNorm(32, in_channels) |
| 313 | self.f_1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) |
| 314 | self.gn_2 = nn.GroupNorm(32, in_channels) |
| 315 | |
| 316 | self.f_2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) |
| 317 | |
| 318 | def forward(self, x, t) -> torch.Tensor: |
| 319 | x_skip = x |