fix

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -238,7 +238,7 @@ def load_vae_weights(
                       pt_list.append("resnets")
                   elif name in ["down_blocks", "up_blocks", "downsamplers", "upsamplers"]:
                       pt_list.append(name)
-                      pt_list.append(idx)
+                      pt_list.append(str(idx))
                   else:
                       pt_list.append(part)
               else:

src/maxdiffusion/tests/test_ltx2_utils.py

@@ -83,7 +83,7 @@ def test_load_transformer_weights(self):
         # But wait, validate_flax_state_dict expects a dict of params, usually just the params subtree.
 
         # We can extract params from state
-        state = nnx.state(model)
+        state = nnx.state(self.transformer)
         # Filter for params? 
         # Usually validate_flax_state_dict expects the full PyTree or a specific dict.
         # state is a State object, acts like a Mapping