| 252 | |
| 253 | |
| 254 | class Upsample(nn.Module): |
| 255 | def __init__(self, dim): |
| 256 | super().__init__() |
| 257 | self.up = nn.Upsample(scale_factor=2, mode="nearest") |
| 258 | self.conv = nn.Conv2d(dim, dim, 3, padding=1) |
| 259 | |
| 260 | def forward(self, x): |
| 261 | return self.conv(self.up(x)) |
| 262 | |
| 263 | |
| 264 | class Downsample(nn.Module): |