diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py
index 1a11a91c75..959ab39577 100644
--- a/tests/studio/run_real_mlx_smoke.py
+++ b/tests/studio/run_real_mlx_smoke.py
@@ -296,13 +296,11 @@ def cmd_train(args) -> int:
             args = config,
         )
 
-        def _on_step(step, total, loss, lr, tok_s, peak_gb, elapsed,
-                     num_tokens, grad_norm = None):
+        def _on_step(
+            step, total, loss, lr, tok_s, peak_gb, elapsed, num_tokens, grad_norm = None
+        ):
             losses_per_step.append(round(float(loss), 4))
-            grad_text = (
-                f"  grad={grad_norm:.4f}"
-                if grad_norm is not None else ""
-            )
+            grad_text = f"  grad={grad_norm:.4f}" if grad_norm is not None else ""
             print(
                 f"  step {step}/{total}  loss={loss:.4f}  lr={lr:.2e}  "
                 f"tok/s={tok_s:.0f}  peak={peak_gb:.2f}GB{grad_text}",