Merge pull request #16683 from balvisio:ba/fix-doc-multi-head-attention

PiperOrigin-RevId: 455435004
keras-team · Jun 16, 2022 · 8930c26 · 8930c26
2 parents 418d246 + c380d05
commit 8930c26
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/keras/layers/attention/multi_head_attention.py b/keras/layers/attention/multi_head_attention.py
@@ -499,8 +499,8 @@ def _compute_attention(
 
         Args:
           query: Projected query `Tensor` of shape `(B, T, N, key_dim)`.
-          key: Projected key `Tensor` of shape `(B, T, N, key_dim)`.
-          value: Projected value `Tensor` of shape `(B, T, N, value_dim)`.
+          key: Projected key `Tensor` of shape `(B, S, N, key_dim)`.
+          value: Projected value `Tensor` of shape `(B, S, N, value_dim)`.
           attention_mask: a boolean mask of shape `(B, T, S)`, that prevents
             attention to certain positions.
           training: Python boolean indicating whether the layer should behave in