diff --git a/mlcolvar/cvs/unsupervised/vae.py b/mlcolvar/cvs/unsupervised/vae.py index 1aaf1148..75420785 100644 --- a/mlcolvar/cvs/unsupervised/vae.py +++ b/mlcolvar/cvs/unsupervised/vae.py @@ -245,6 +245,7 @@ def training_step(self, train_batch, batch_idx): # Log. name = "train" if self.training else "valid" self.log(f"{name}_loss", loss, on_epoch=True) + self.log("lr", self.optimizers().param_groups[0]['lr'], on_epoch=True) self.log(f"{name}_reconstruction_loss", reconstruction_loss, on_epoch=True) self.log(f"{name}_kl_loss", kl_loss, on_epoch=True) self.log(f"beta", self.beta, on_epoch=True)