From bb8227f18197202e3196ff7fbb1f83e329df7768 Mon Sep 17 00:00:00 2001 From: Yi Wang Date: Mon, 10 Jun 2024 14:47:31 -0700 Subject: [PATCH] Correct type annotation of llama.ModelArgs.num_key_value_heads (#827) --- llms/mlx_lm/models/llama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llms/mlx_lm/models/llama.py b/llms/mlx_lm/models/llama.py index 55a2b5db..e7f4f16a 100644 --- a/llms/mlx_lm/models/llama.py +++ b/llms/mlx_lm/models/llama.py @@ -16,7 +16,7 @@ class ModelArgs(BaseModelArgs): num_attention_heads: int rms_norm_eps: float vocab_size: int - num_key_value_heads: int = None + num_key_value_heads: Optional[int] = None attention_bias: bool = False mlp_bias: bool = False rope_theta: float = 10000