Merge pull request #1876 from huggingface/mean-fix

Mean does not exist in TF2
2019-11-29 09:26:33 +01:00
parent f0d22b6363 3de31f8d28
commit 1ab8dc44b3
1 changed files with 1 additions and 1 deletions
--- a/transformers/modeling_tf_utils.py
+++ b/transformers/modeling_tf_utils.py
@@ -460,7 +460,7 @@ class TFSequenceSummary(tf.keras.layers.Layer):
        elif self.summary_type == 'first':
            output = hidden_states[:, 0]
        elif self.summary_type == 'mean':
-            output = tf.mean(hidden_states, axis=1)
+            output = tf.reduce_mean(hidden_states, axis=1)
        elif self.summary_type == 'cls_index':
            hidden_shape = shape_list(hidden_states)  # e.g. [batch, num choices, seq length, hidden dims]
            if cls_index is None: