new fix

Author: prishajain1

src/maxdiffusion/tests/ltx2_vae_parity_test.py

@@ -122,7 +122,7 @@ def fix_keys(d):
 
     # Also Check Encoder Latents
     print("\nEncoder Latents Stats:")
-    posterior = model.encode(jax_input)
+    posterior = model.encode(jax_input).latent_dist
     # posterior is DiagonalGaussianDistribution
     # Check mode
     latents = posterior.mode()