| 41 | return [result, mu, log_var] |
| 42 | |
| 43 | class LadderBlock(nn.Module): |
| 44 | def __init__(self, |
| 45 | in_channels: int, |
| 46 | latent_dim: int): |
| 47 | super(LadderBlock, self).__init__() |
| 48 | |
| 49 | # Build Decoder |
| 50 | self.decode = nn.Sequential(nn.Linear(in_channels, latent_dim), |
| 51 | nn.BatchNorm1d(latent_dim)) |
| 52 | self.fc_mu = nn.Linear(latent_dim, latent_dim) |
| 53 | self.fc_var = nn.Linear(latent_dim, latent_dim) |
| 54 | |
| 55 | def forward(self, z: Tensor) -> Tensor: |
| 56 | z = self.decode(z) |
| 57 | mu = self.fc_mu(z) |
| 58 | log_var = self.fc_var(z) |
| 59 | |
| 60 | return [mu, log_var] |
| 61 | |
| 62 | class LVAE(BaseVAE): |
| 63 | |