fix

prishajain1 · prishajain1 · commit 42a2c64f3116 · 2026-02-24T15:45:43.000+05:30
diff --git a/src/maxdiffusion/tests/test_ltx2_utils.py b/src/maxdiffusion/tests/test_ltx2_utils.py
@@ -7,7 +7,7 @@
 from maxdiffusion.models.ltx2.autoencoder_kl_ltx2 import LTX2VideoAutoencoderKL
 from maxdiffusion.models.ltx2.ltx2_utils import load_transformer_weights, load_vae_weights
 from maxdiffusion.models.modeling_flax_pytorch_utils import validate_flax_state_dict
-from flax.traverse_util import flatten_dict
+from flax.traverse_util import flatten_dict, unflatten_dict
 
 class LTX2VideoConfig:
     def __init__(self):
