Internal change

PiperOrigin-RevId: 347490358

official/nlp/configs/encoders.py

@@ -26,7 +26,6 @@ import tensorflow as tf
 
 from official.modeling import hyperparams
 from official.modeling import tf_utils
-from official.nlp import keras_nlp
 from official.nlp.modeling import networks
 from official.nlp.projects.bigbird import encoder as bigbird_encoder
 
@@ -186,9 +185,8 @@ ENCODER_CLS = {
 
 
 @gin.configurable
-def build_encoder(
-    config: EncoderConfig,
-    embedding_layer: Optional[keras_nlp.layers.OnDeviceEmbedding] = None,
+def build_encoder(config: EncoderConfig,
+                  embedding_layer: Optional[tf.keras.layers.Layer] = None,
                   encoder_cls=None,
                   bypass_config: bool = False):
   """Instantiate a Transformer encoder network from EncoderConfig.