Fix TransfoXL (#9302)

2020-12-28 20:52:18 +01:00
parent d97d06d05f
commit 64103fb6be
1 changed files with 2 additions and 2 deletions
--- a/src/transformers/models/transfo_xl/modeling_tf_transfo_xl.py
+++ b/src/transformers/models/transfo_xl/modeling_tf_transfo_xl.py
@@ -501,8 +501,8 @@ class TFTransfoXLMainLayer(tf.keras.layers.Layer):

        # There are `mlen + qlen` steps that can be cached into mems
        new_mems = []
-        end_idx = mlen + max(0, qlen)
-        beg_idx = max(0, end_idx - self.mem_len)
+        end_idx = mlen + tf.math.maximum(0, qlen)
+        beg_idx = tf.math.maximum(0, end_idx - tf.convert_to_tensor(self.mem_len))
        for i in range(len(hids)):

            cat = tf.concat([mems[i], hids[i]], axis=0)