convert Dimension to int to avoid "TypeError: unsupported operand type(s) for...

convert Dimension to int to avoid "TypeError: unsupported operand type(s) for /: float and Dimension"

PiperOrigin-RevId: 358988747

official/nlp/modeling/layers/transformer.py

@@ -202,7 +202,7 @@ class TransformerDecoderBlock(tf.keras.layers.Layer):
       raise ValueError(
           "The hidden size (%d) is not a multiple of the number of attention "
           "heads (%d)" % (hidden_size, self.num_attention_heads))
-    self.attention_head_size = int(hidden_size / self.num_attention_heads)
+    self.attention_head_size = int(hidden_size) // self.num_attention_heads
     common_kwargs = dict(
         bias_initializer=self._bias_initializer,
         kernel_regularizer=self._kernel_regularizer,

official/nlp/modeling/models/seq2seq_transformer.py

@@ -591,5 +591,6 @@ class TransformerDecoder(tf.keras.layers.Layer):
 
 def attention_initializer(hidden_size):
   """Initializer for attention layers in Seq2SeqTransformer."""
+  hidden_size = int(hidden_size)
   limit = math.sqrt(6.0 / (hidden_size + hidden_size))
   return tf.keras.initializers.RandomUniform(minval=-limit, maxval=limit)