Use kl_coeff in loss function calculation (#330)

* Update basic_ae_module.py remove unnecessary kl_coeff argument * included kl_coeff in the calculation of the loss function
Lightning-Universe · Nov 2, 2020 · d4e6096 · d4e6096
1 parent 978fa1c
commit d4e6096
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 1 deletion.
diff --git a/pl_bolts/models/autoencoders/basic_ae/basic_ae_module.py b/pl_bolts/models/autoencoders/basic_ae/basic_ae_module.py
@@ -36,7 +36,6 @@ def __init__(
         first_conv: bool = False,
         maxpool1: bool = False,
         enc_out_dim: int = 512,
-        kl_coeff: float = 0.1,
         latent_dim: int = 256,
         lr: float = 1e-4,
         **kwargs

diff --git a/pl_bolts/models/autoencoders/basic_vae/basic_vae_module.py b/pl_bolts/models/autoencoders/basic_vae/basic_vae_module.py
@@ -127,6 +127,7 @@ def step(self, batch, batch_idx):
 
         kl = log_qz - log_pz
         kl = kl.mean()
+        kl *= self.kl_coeff
 
         loss = kl + recon_loss