new fix

Author: prishajain1

src/maxdiffusion/tests/ltx2_vae_parity_test.py

@@ -110,7 +110,7 @@ def fix_keys(d):
     print("Running Flax forward pass...")
     # Call the model
     # Note: default deterministic=True, causal=True/False depending on init
-    jax_recon = model(jax_input, sample_posterior=False)
+    jax_recon = model(jax_input, sample_posterior=False).sample
 
     # 5. Print Output Stats
     print("\nOutput Stats:")