fix

prishajain1 · prishajain1 · commit 5fd61964f0f5 · 2026-02-23T11:02:45.000+05:30
diff --git a/src/maxdiffusion/tests/test_ltx2_utils.py b/src/maxdiffusion/tests/test_ltx2_utils.py
@@ -122,7 +122,8 @@ def test_load_vae_weights(self):
                 continue
             filtered_eval_shapes[k] = v
             
-        validate_flax_state_dict(filtered_eval_shapes, flatten_dict(loaded_weights))
+        from flax.traverse_util import unflatten_dict
+        validate_flax_state_dict(unflatten_dict(filtered_eval_shapes), flatten_dict(loaded_weights))
         print("VAE Weights Validated Successfully!")
 
 if __name__ == "__main__":
