diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py
index bf533e353c..d6f6eb329c 100644
--- a/src/maxtext/kernels/megablox/ops.py
+++ b/src/maxtext/kernels/megablox/ops.py
@@ -91,7 +91,10 @@ def gmm(
     # 2. for batchsplit, explicitly pass the rule
     quantization_rule = qwix_rule if qwix_rule else qpl.get_current_rule("gmm")
     if not quantization_rule or not isinstance(quantization_rule, qwix.QtRule):
-      raise ValueError("Expect a QtRule for quantized training. " f"But get quantization_rule={quantization_rule}")
+      raise ValueError(
+          "Expect a QtRule for quantized training. "
+          f"But get quantization_rule={quantization_rule}"
+      )
   else:
     # Handcraft a rule that matches the AQT's behavior.
     if lhs_quantize_dtype or rhs_quantize_dtype: