Unverified Commit 96cc02b5 authored by yis11178's avatar yis11178 Committed by GitHub
Browse files

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


Co-authored-by: default avataryis <yis@graphcore.ai>
parent 43f953cc
......@@ -170,7 +170,7 @@ class TFMultiHeadSelfAttention(tf.keras.layers.Layer):
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]
......
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