debug

Author: prishajain1

src/maxdiffusion/models/ltx_2/transformer_ltx2.py

@@ -263,7 +263,13 @@ def __call__(
         # table shape: (6, dim) -> (1, 1, 6, dim)
         scale_shift_table_reshaped = jnp.expand_dims(self.scale_shift_table, axis=(0, 1))
         # temb shape: (batch, temb_dim) -> (batch, 1, 6, dim)  (assuming temb_dim is num_ada_params * dim)
+        print(f"DEBUG: scale_shift_table_reshaped shape: {scale_shift_table_reshaped.shape}")
+        print(f"DEBUG: temb shape before reshape: {temb.shape}")
+        
         temb_reshaped = temb.reshape(batch_size, 1, num_ada_params, -1)
+        
+        print(f"DEBUG: temb_reshaped shape: {temb_reshaped.shape}")
+        
         ada_values = scale_shift_table_reshaped + temb_reshaped
 
         shift_msa = ada_values[:, :, 0, :]