From a1a67a3ced35531e871026a01cfff653105d532f Mon Sep 17 00:00:00 2001 From: abhishek thakur Date: Wed, 3 Feb 2021 10:49:07 +0100 Subject: [PATCH] Fix GroupedLinearLayer in TF ConvBERT (#9972) --- src/transformers/models/convbert/modeling_tf_convbert.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/transformers/models/convbert/modeling_tf_convbert.py b/src/transformers/models/convbert/modeling_tf_convbert.py index fbd9ecb4f4..59e0c3362b 100644 --- a/src/transformers/models/convbert/modeling_tf_convbert.py +++ b/src/transformers/models/convbert/modeling_tf_convbert.py @@ -435,9 +435,10 @@ class GroupedLinearLayer(tf.keras.layers.Layer): ) def call(self, hidden_states): - batch_size = shape_list(tensor=hidden_states)[1] - x = tf.reshape(hidden_states, [-1, self.num_groups, self.group_in_dim]) - x = tf.matmul(a=x, b=self.kernel, transpose_b=True) + batch_size = shape_list(hidden_states)[0] + x = tf.transpose(tf.reshape(hidden_states, [-1, self.num_groups, self.group_in_dim]), [1, 0, 2]) + x = tf.matmul(x, self.kernel) + x = tf.transpose(x, [1, 0, 2]) x = tf.reshape(x, [batch_size, -1, self.output_size]) x = tf.nn.bias_add(value=x, bias=self.bias) return x