pipeline fix

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1359,7 +1359,7 @@ def __call__(
 
     # If weights are on CPU (from a previous offload), restore them to TPU with correct sharding
     sample_leaf = jax.tree_util.tree_leaves(state)[0]
-    if sample_leaf.device().platform == "cpu":
+    if list(sample_leaf.devices())[0].platform == "cpu":
       logical_state_spec = nnx.get_partition_spec(state)
       logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, self.mesh, self.config.logical_axis_rules)
       state = jax.tree_util.tree_map(lambda x, sharding: jax.device_put(x, sharding), state, logical_state_sharding)