From 69cc78ef535f949b4d0e1746408c40b04c45c713 Mon Sep 17 00:00:00 2001 From: Maitry Sinha Date: Thu, 26 Feb 2026 15:05:59 +0530 Subject: [PATCH 1/2] fix: correct total_loss calculate assign --- examples/generative/vae.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/generative/vae.py b/examples/generative/vae.py index ce7f1ca086..34c03ceacb 100644 --- a/examples/generative/vae.py +++ b/examples/generative/vae.py @@ -110,13 +110,13 @@ def train_step(self, data): total_loss = reconstruction_loss + kl_loss grads = tape.gradient(total_loss, self.trainable_weights) self.optimizer.apply_gradients(zip(grads, self.trainable_weights)) - self.total_loss_tracker.update_state(total_loss) self.reconstruction_loss_tracker.update_state(reconstruction_loss) self.kl_loss_tracker.update_state(kl_loss) + self.total_loss_tracker.update_state(total_loss) return { - "loss": self.total_loss_tracker.result(), "reconstruction_loss": self.reconstruction_loss_tracker.result(), "kl_loss": self.kl_loss_tracker.result(), + "loss": self.total_loss_tracker.result(), } From b2ac210b9db2c4ba7e34423a694e961055c8dd58 Mon Sep 17 00:00:00 2001 From: Maitry Sinha Date: Thu, 26 Feb 2026 15:21:21 +0530 Subject: [PATCH 2/2] Update examples/generative/vae.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- examples/generative/vae.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/generative/vae.py b/examples/generative/vae.py index 34c03ceacb..091e2ef306 100644 --- a/examples/generative/vae.py +++ b/examples/generative/vae.py @@ -116,7 +116,7 @@ def train_step(self, data): return { "reconstruction_loss": self.reconstruction_loss_tracker.result(), "kl_loss": self.kl_loss_tracker.result(), - "loss": self.total_loss_tracker.result(), + "total_loss": self.total_loss_tracker.result(), }