fix

Author: susanbao

src/maxdiffusion/trainers/wan_trainer.py

@@ -480,4 +480,4 @@ def loss_fn(params):
   metrics = {"scalar": {"learning/eval_loss": loss}}
 
   # Return the computed metrics and the new RNG key for the next eval step
-  return metrics, new_rng, 
+  return metrics, new_rng