fix formating

AI-Hypercomputer · Feb 19, 2025 · 8b6d384 · 8b6d384
1 parent 66d68e3
commit 8b6d384
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/MaxText/layers/attentions.py b/MaxText/layers/attentions.py
@@ -1386,7 +1386,9 @@ def setup(self):
     super().setup()
 
     # Assert required configuration parameters for MLA attention.
-    assert self.config.attention_type == AttentionType.MLA.value, f"MLA requires MLA attention type {AttentionType.MLA.value}"
+    assert (
+        self.config.attention_type == AttentionType.MLA.value
+    ), f"MLA requires MLA attention type {AttentionType.MLA.value}"
     assert self.kv_lora_rank > 0, "KV LoRA rank must be > 0"
     assert self.qk_nope_head_dim > 0, "QK NoPe head dim must be > 0"
     assert self.qk_rope_head_dim > 0, "QK RoPE head dim must be > 0"