| 494 | |
| 495 | |
| 496 | class Downsample2D(nn.Module): |
| 497 | |
| 498 | def __init__(self, channels): |
| 499 | super().__init__() |
| 500 | self.conv = nn.Conv2d(channels, channels, 3, stride=2, padding=1) |
| 501 | |
| 502 | def forward(self, x): |
| 503 | return self.conv(x) |
| 504 | |
| 505 | |
| 506 | class SpatialTransformer(nn.Module): |