No public description

PiperOrigin-RevId: 670736074

Author: joezoug

official/nlp/modeling/networks/mobile_bert_encoder.py

@@ -162,7 +162,10 @@ def __init__(self,
         pooled_output=first_token,
         encoder_outputs=all_layer_outputs,
         attention_scores=all_attention_scores)
+    super().__init__(
+        inputs=self.inputs, outputs=outputs, **kwargs)
     self._config = dict(
+        name=self.name,
         word_vocab_size=word_vocab_size,
         word_embed_size=word_embed_size,
         type_vocab_size=type_vocab_size,
@@ -182,11 +185,9 @@ def __init__(self,
         normalization_type=normalization_type,
         classifier_activation=classifier_activation,
         input_mask_dtype=input_mask_dtype,
+        **kwargs,
     )
 
-    super().__init__(
-        inputs=self.inputs, outputs=outputs, **kwargs)
-
   def get_config(self):
     return dict(self._config)