(self, params: AutoEncoderParams)
| 276 | |
| 277 | class AutoEncoder(nn.Module): |
| 278 | def __init__(self, params: AutoEncoderParams): |
| 279 | super().__init__() |
| 280 | self.encoder = Encoder( |
| 281 | resolution=params.resolution, |
| 282 | in_channels=params.in_channels, |
| 283 | ch=params.ch, |
| 284 | ch_mult=params.ch_mult, |
| 285 | num_res_blocks=params.num_res_blocks, |
| 286 | z_channels=params.z_channels, |
| 287 | ) |
| 288 | self.decoder = Decoder( |
| 289 | resolution=params.resolution, |
| 290 | in_channels=params.in_channels, |
| 291 | ch=params.ch, |
| 292 | out_ch=params.out_ch, |
| 293 | ch_mult=params.ch_mult, |
| 294 | num_res_blocks=params.num_res_blocks, |
| 295 | z_channels=params.z_channels, |
| 296 | ) |
| 297 | self.reg = DiagonalGaussian() |
| 298 | |
| 299 | self.scale_factor = params.scale_factor |
| 300 | self.shift_factor = params.shift_factor |
| 301 | |
| 302 | def encode(self, x: Tensor) -> Tensor: |
| 303 | z = self.reg(self.encoder(x)) |
nothing calls this directly
no test coverage detected