fix in transformer

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1009,29 +1009,29 @@ def __call__(
     with self.conditional_named_scope("transformer_block"):
       def scan_fn(carry, block):
         hidden_states, audio_hidden_states, rngs_carry = carry
-      hidden_states_out, audio_hidden_states_out = block(
-          hidden_states=hidden_states,
-          audio_hidden_states=audio_hidden_states,
-          encoder_hidden_states=encoder_hidden_states,
-          audio_encoder_hidden_states=audio_encoder_hidden_states,
-          temb=temb,
-          temb_audio=temb_audio,
-          temb_ca_scale_shift=video_cross_attn_scale_shift,
-          temb_ca_audio_scale_shift=audio_cross_attn_scale_shift,
-          temb_ca_gate=video_cross_attn_a2v_gate,
-          temb_ca_audio_gate=audio_cross_attn_v2a_gate,
-          video_rotary_emb=video_rotary_emb,
-          audio_rotary_emb=audio_rotary_emb,
-          ca_video_rotary_emb=video_cross_attn_rotary_emb,
-          ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
-          encoder_attention_mask=encoder_attention_mask,
-          audio_encoder_attention_mask=audio_encoder_attention_mask,
-      )
-      return (
-          hidden_states_out.astype(hidden_states.dtype),
-          audio_hidden_states_out.astype(audio_hidden_states.dtype),
-          rngs_carry,
-      ), None
+        hidden_states_out, audio_hidden_states_out = block(
+            hidden_states=hidden_states,
+            audio_hidden_states=audio_hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            audio_encoder_hidden_states=audio_encoder_hidden_states,
+            temb=temb,
+            temb_audio=temb_audio,
+            temb_ca_scale_shift=video_cross_attn_scale_shift,
+            temb_ca_audio_scale_shift=audio_cross_attn_scale_shift,
+            temb_ca_gate=video_cross_attn_a2v_gate,
+            temb_ca_audio_gate=audio_cross_attn_v2a_gate,
+            video_rotary_emb=video_rotary_emb,
+            audio_rotary_emb=audio_rotary_emb,
+            ca_video_rotary_emb=video_cross_attn_rotary_emb,
+            ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
+            encoder_attention_mask=encoder_attention_mask,
+            audio_encoder_attention_mask=audio_encoder_attention_mask,
+        )
+        return (
+            hidden_states_out.astype(hidden_states.dtype),
+            audio_hidden_states_out.astype(audio_hidden_states.dtype),
+            rngs_carry,
+        ), None
 
     if self.scan_layers:
       rematted_scan_fn = self.gradient_checkpoint.apply(