From 8a265f0d540ae797bb7d32656605d6ec38938ec4 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 24 Apr 2024 18:52:43 -0400 Subject: [PATCH] Fix incorrect type annotation (#720) A `Tuple` is missing in this type annotation. --- llms/speculative_decoding/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llms/speculative_decoding/model.py b/llms/speculative_decoding/model.py index 2300fb6c..c310b943 100644 --- a/llms/speculative_decoding/model.py +++ b/llms/speculative_decoding/model.py @@ -105,7 +105,7 @@ class MultiHeadAttention(nn.Module): values: mx.array, mask: Optional[mx.array], cache: Optional[Tuple[mx.array, mx.array]] = None, - ) -> [mx.array, Tuple[mx.array, mx.array]]: + ) -> Tuple[mx.array, Tuple[mx.array, mx.array]]: queries = self.query_proj(queries) keys = self.key_proj(keys) values = self.value_proj(values)