| 256 | return ret |
| 257 | |
| 258 | class Downsample: |
| 259 | def __init__(self, channels): |
| 260 | self.op = Conv2d(channels, channels, 3, stride=2, padding=1) |
| 261 | |
| 262 | def __call__(self, x): |
| 263 | return self.op(x) |
| 264 | |
| 265 | class Upsample: |
| 266 | def __init__(self, channels): |