Skip to content

GitLab

Commit 1c507995 authored by thomwolf's avatar thomwolf
Browse files

simpler distilbert mask - fix tf tests

parent 58b302ca
Show whitespace changes
Inline Side-by-side
Showing with 0 additions and 5 deletions
transformers/modeling_distilbert.py
View file @ 1c507995
...@@ -159,8 +159,6 @@ class MultiHeadSelfAttention(nn.Module): ...@@ -159,8 +159,6 @@ class MultiHeadSelfAttention(nn.Module):
dim_per_head = self.dim // self.n_heads dim_per_head = self.dim // self.n_heads
assert 2 <= mask.dim() <= 3
causal = (mask.dim() == 3)
mask_reshp = (bs, 1, 1, k_length) mask_reshp = (bs, 1, 1, k_length)
def shape(x): def shape(x):
......
transformers/modeling_tf_distilbert.py
View file @ 1c507995
...@@ -226,9 +226,6 @@ class TFMultiHeadSelfAttention(tf.keras.layers.Layer): ...@@ -226,9 +226,6 @@ class TFMultiHeadSelfAttention(tf.keras.layers.Layer):
dim_per_head = self.dim // self.n_heads dim_per_head = self.dim // self.n_heads
mask_shape = shape_list(mask)
assert 2 <= len(mask_shape) <= 3
causal = (mask_shape) == 3)
mask_reshape = [bs, 1, 1, k_length] mask_reshape = [bs, 1, 1, k_length]
def shape(x): def shape(x):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment