diff --git a/Matrix-Game-2/utils/wan_wrapper.py b/Matrix-Game-2/utils/wan_wrapper.py
index 537a1f4..ccd3976 100644
--- a/Matrix-Game-2/utils/wan_wrapper.py
+++ b/Matrix-Game-2/utils/wan_wrapper.py
@@ -58,7 +58,7 @@ def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torc
             decode_function = self.model.decode
 
         output = []
-        for u in zs:
+        for u in latent:
             output.append(decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0))
         output = torch.stack(output, dim=0)
         return output
@@ -206,4 +206,3 @@ def post_init(self):
         We can gradually add more methods here if needed.
         """
         self.get_scheduler()
-