| 70 | |
| 71 | |
| 72 | class Upsample(nn.Module): |
| 73 | def __init__(self, dim, save_fm=False): |
| 74 | super().__init__() |
| 75 | self.up = nn.Upsample(scale_factor=2, mode="nearest") |
| 76 | self.conv = nn.Conv2d(dim, dim, 3, padding=1) |
| 77 | self.save_fm = save_fm |
| 78 | |
| 79 | def forward(self, x): |
| 80 | x = self.conv(self.up(x)) |
| 81 | return x |
| 82 | |
| 83 | |
| 84 | class Downsample(nn.Module): |