Convert NumPy attention mask to TensorFlow tensor for shape_list

2020-02-04 15:19:22 -05:00
1 changed files with 2 additions and 0 deletions
@@ -365,6 +365,8 @@ class TFXLMMainLayer(tf.keras.layers.Layer):
        #     assert src_enc.size(0) == bs

        # generate masks
+        if attention_mask is not None:
+            attention_mask = tf.cast(attention_mask, tf.int32)
        mask, attn_mask = get_masks(slen, lengths, self.causal, padding_mask=attention_mask)
        # if self.is_decoder and src_enc is not None:
        #     src_mask = torch.arange(src_len.max(), dtype=torch.long, device=lengths.device) < src_len[:, None]