converting timesteps to jax array

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1237,7 +1237,9 @@ def run_connectors(graphdef, state, hidden_states, attention_mask):
       )
 
       import time
-      for i, t in enumerate(timesteps):
+      timesteps_jax = jnp.array(timesteps, dtype=jnp.float32)
+      for i, t_val in enumerate(timesteps):
+        t = timesteps_jax[i]
 
         # Isolate input sharding to scan_layers=False to avoid affecting the standard path
         latents_jax_sharded = latents_jax