| 24 | self.kl = (sigma**2 + mu**2 - torch.log(sigma) - 1/2).sum() |
| 25 | return z |
| 26 | class Decoder(nn.Module): |
| 27 | def __init__(self,latent_dims): |
| 28 | super(Decoder,self).__init__() |
| 29 | self.linear1 = nn.Linear(latent_dims,64) |
| 30 | self.linear2 = nn.Linear(64,128) |
| 31 | self.linear3 = nn.Linear(128,256) |
| 32 | |
| 33 | def forward(self,z): |
| 34 | z = F.relu(self.linear1(z)) |
| 35 | z = F.relu(self.linear2(z)) |
| 36 | z = self.linear3(z) |
| 37 | return z |
| 38 | |
| 39 | class VariationalAutoencoder(nn.Module): |
| 40 | def __init__(self, latent_dims): |