Nit fix in attention

keras-team · Apr 25, 2024 · 11d3314 · 11d3314
1 parent ac94325
commit 11d3314
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/keras_nlp/models/phi3/phi3_attention.py b/keras_nlp/models/phi3/phi3_attention.py
@@ -70,7 +70,7 @@ def build(self, inputs_shape):
         self._query_dense.build(inputs_shape)
 
         self._key_dense = keras.layers.Dense(
-            self.head_dim * self.num_key_value_groups,
+            self.head_dim * self.num_key_value_heads,
             kernel_initializer=self.kernel_initializer,
             use_bias=False,
             dtype=self.dtype_policy,
@@ -79,7 +79,7 @@ def build(self, inputs_shape):
         self._key_dense.build(inputs_shape)
 
         self._value_dense = keras.layers.Dense(
-            self.head_dim * self.num_key_value_groups,
+            self.head_dim * self.num_key_value_heads,
             kernel_initializer=self.kernel_initializer,
             use_bias=False,
             dtype=self.dtype_policy,