[bugfix] fix bz >1 for training (#477)

Author: SolitaryThinker

fastvideo/v1/training/wan_training_pipeline.py

@@ -147,7 +147,7 @@ def train_one_step(
                         dtype=torch.bfloat16)
                 with set_forward_context(current_timestep=timesteps,
                                          attn_metadata=None):
-                    model_pred = transformer(**input_kwargs)[0]
+                    model_pred = transformer(**input_kwargs)
 
                 if precondition_outputs:
                     model_pred = noisy_model_input - model_pred * sigmas