From 84cf78361641bf3bae71a37513383995eb314e53 Mon Sep 17 00:00:00 2001 From: Yunpeng Li Date: Thu, 10 Nov 2022 10:53:18 -0800 Subject: [PATCH] Internal change PiperOrigin-RevId: 487583744 --- official/nlp/modeling/layers/kernel_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/official/nlp/modeling/layers/kernel_attention.py b/official/nlp/modeling/layers/kernel_attention.py index cbfd6ba7b..76e6b123b 100644 --- a/official/nlp/modeling/layers/kernel_attention.py +++ b/official/nlp/modeling/layers/kernel_attention.py @@ -252,7 +252,7 @@ def causal_windowed_performer_attention(query_matrix, attention = numerator / denominator attention = tf.reshape(attention, new_shape) - start = tf.zeros([len(old_shape)], dtype=old_shape.dtype) + start = tf.zeros([old_shape.shape[0]], dtype=old_shape.dtype) attention = tf.slice(attention, start, old_shape) # Queued window cache (drop instead of decay) not yet supported. -- GitLab