debug

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -390,6 +390,7 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
           sharding = NamedSharding(mesh, P())
         state[path].value = device_put_replicated(val, sharding)
       else:
+        max_logging.log(f"[Tuning] VAE Key {path} fell into default sharding fallback!")
         state[path].value = jax.device_put(val)
 
     state = nnx.from_flat_state(state)