On the VAE model, `kl_coeff` is never applied in the loss function

[miccio-dk]

🐛 Bug

The VAE model contained here takes an optional kl_coeff parameter that's supposed to be a scaling factor for the KL term of the variational loss. This might be useful to avoid the "posterior collapse phenomenon". However, while inspecting the codebase (i.e. here) I can see that the parameter is initialized, stored, but never used.

To Reproduce

Steps to reproduce the behavior:

1. Train a VAE
2. Set a custom value for the kl_coeff argument
3. Train again

Code sample

Code for achieving this can be taken from https://github.com/PyTorchLightning/pytorch-lightning-bolts/blob/master/tests/models/test_autoencoders.py#L13 but the behavior should be obvious from basic_vae_module.py

Expected behavior

The loss function calculation (https://github.com/PyTorchLightning/pytorch-lightning-bolts/blob/master/pl_bolts/models/autoencoders/basic_vae/basic_vae_module.py#L119) should account for kl_loss.

Environment

• PyTorch Version (e.g., 1.0): 1.8 (nightly)
• OS (e.g., Linux): Linux
• How you installed PyTorch (conda, pip, source): conda
• Build command you used (if compiling from source): N/A
• Python version: 3.8
• CUDA/cuDNN version: N/A
• GPU models and configuration: N/A
• Any other relevant information: N/A

Additional context

Thanks for the amazing library/libraries! :)

[github-actions[bot]]

Hi! thanks for your contribution!, great first issue!

[akihironitta]

@miccio-dk Thanks for reporting the issue! Good catch. kl_coeff is actually never used. Would you be interested in submitting a PR to fix this?

[miccio-dk]

Absolutely! I suppose i just have to multiply it with kl in https://github.com/PyTorchLightning/pytorch-lightning-bolts/blob/master/pl_bolts/models/autoencoders/basic_vae/basic_vae_module.py#L131. I'll also remove the kl_coeff argument from the basic autoencoder, where it was probably added by mistake. :)

[ananyahjha93]

@miccio-dk yes! these 2 things should be enough for the PR. Thanks!

[akihironitta]

@miccio-dk Sounds nice! Thanks :]

[miccio-dk]

Aaaaand here it is!