Fix activation hinting.

PiperOrigin-RevId: 414032298

official/nlp/modeling/networks/bert_encoder.py

@@ -23,6 +23,8 @@ from official.nlp.modeling import layers
 
 
 _Initializer = Union[str, tf.keras.initializers.Initializer]
+_Activation = Union[str, Callable[..., Any]]
+
 _approx_gelu = lambda x: tf.keras.activations.gelu(x, approximate=True)
 
 
@@ -83,7 +85,7 @@ class BertEncoderV2(tf.keras.layers.Layer):
       max_sequence_length: int = 512,
       type_vocab_size: int = 16,
       inner_dim: int = 3072,
-      inner_activation: Callable[..., Any] = _approx_gelu,
+      inner_activation: _Activation = _approx_gelu,
       output_dropout: float = 0.1,
       attention_dropout: float = 0.1,
       initializer: _Initializer = tf.keras.initializers.TruncatedNormal(