From 39c318ee216726fc704a0b9656f3fc8b8d05d52d Mon Sep 17 00:00:00 2001
From: HGB <87762857+radna0@users.noreply.github.com>
Date: Mon, 18 Nov 2024 00:21:14 +0700
Subject: [PATCH] Update latent_preview.py

---
 latent_preview.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/latent_preview.py b/latent_preview.py
index ee1344128..6ed14d3a8 100644
--- a/latent_preview.py
+++ b/latent_preview.py
@@ -144,9 +144,12 @@ def prepare_callback(model, steps, x0_output_dict=None):
         if previewer:
             preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
         pbar.update_absolute(step + 1, total_steps, preview_bytes)
-        if args.xla or args.xla_fsdp:
+
+        if args.xla or args.xla_spmd:
             import torch_xla as xla
 
             xla.sync()
 
+        pbar.update_absolute(step + 1, total_steps, preview_bytes)
+
     return callback