Why VAE model in pytorch dosn't use torch.nn.KLDivLoss?

  • I saw several examples of VAE implementations.
  • In VAE we use loss with 2 parts: MSE and KLDivLoss
  • In all the examples I saw, they wrote the VAE loss (MSE+KL LOSS) in their own implementation and didn’t used: torch.nn.KLDivLoss

One example can be found here:
https://github.com/AntixK/PyTorch-VAE/blob/master/models/vanilla_vae.py

Why did they implement the part of the KLDivLoss and didn’t use torch.nn.KLDivLoss ?

1 Like