Internal change

PiperOrigin-RevId: 310981542

Internal change
PiperOrigin-RevId: 310981542
36d27f01 · A. Unique TensorFlower · ebc3edc6 · 36d27f01
Commit 36d27f01 authored May 11, 2020 by A. Unique TensorFlower
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

official/nlp/transformer/attention_layer.py official/nlp/transformer/attention_layer.py +1 -1

No files found.
--- a/official/nlp/transformer/attention_layer.py
+++ b/official/nlp/transformer/attention_layer.py
@@ -54,7 +54,7 @@ class Attention(tf.keras.layers.Layer):
      limit = math.sqrt(6.0 / (fan_in + fan_out))
      return tf.keras.initializers.RandomUniform(minval=-limit, maxval=limit)

-    attention_initializer = _glorot_initializer(input_shape[-1],
+    attention_initializer = _glorot_initializer(input_shape.as_list()[-1],
                                                self.hidden_size)
    self.query_dense_layer = layers.DenseEinsum(
        output_shape=(self.num_heads, size_per_head),