### Scale up the feature map by $2 \times$
| 175 | |
| 176 | |
| 177 | class Upsample(nn.Module): |
| 178 | """ |
| 179 | ### Scale up the feature map by $2 \times$ |
| 180 | """ |
| 181 | |
| 182 | def __init__(self, n_channels): |
| 183 | super().__init__() |
| 184 | self.conv = nn.ConvTranspose2d(n_channels, n_channels, (4, 4), (2, 2), (1, 1)) |
| 185 | |
| 186 | def forward(self, x: torch.Tensor, t: torch.Tensor): |
| 187 | # `t` is not used, but it's kept in the arguments because for the attention layer function signature |
| 188 | # to match with `ResidualBlock`. |
| 189 | _ = t |
| 190 | return self.conv(x) |
| 191 | |
| 192 | |
| 193 | class Downsample(nn.Module): |