change tf.math.divide with int(/) to remove dim_per_head from the TF …

…graph (#14600)

Co-authored-by: yis <yis@graphcore.ai>

src/transformers/models/distilbert/modeling_tf_distilbert.py

@@ -170,7 +170,7 @@ def call(self, query, key, value, mask, head_mask, output_attentions, training=F
         k_length = shape_list(key)[1]
         # assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured'
         # assert key.size() == value.size()
-        dim_per_head = tf.math.divide(self.dim, self.n_heads)
+        dim_per_head = int(self.dim / self.n_heads)
         dim_per_head = tf.cast(dim_per_head, dtype=tf.int32)
         mask_reshape = [bs, 1, 1, k_length]