diff --git a/official/nlp/modeling/networks/encoder_scaffold.py b/official/nlp/modeling/networks/encoder_scaffold.py
index 68dad751d0444e32c10a35b96edc245ef8e0ba6f..186512357f7eb187c7d5a578114e16719c0a99cc 100644
--- a/official/nlp/modeling/networks/encoder_scaffold.py
+++ b/official/nlp/modeling/networks/encoder_scaffold.py
@@ -114,6 +114,9 @@ class EncoderScaffold(tf.keras.Model):
       inputs = embedding_network.inputs
       embeddings, attention_mask = embedding_network(inputs)
       embedding_layer = None
+      position_embedding_layer = None
+      type_embedding_layer = None
+      embedding_norm_layer = None
     else:
       embedding_network = None
       seq_length = embedding_cfg.get('seq_length', None)
@@ -224,6 +227,10 @@ class EncoderScaffold(tf.keras.Model):
 
     self._embedding_layer = embedding_layer
     self._embedding_network = embedding_network
+    self._position_embedding_layer = position_embedding_layer
+    self._type_embedding_layer = type_embedding_layer
+    self._embedding_norm_layer = embedding_norm_layer
+    self._embedding_network = embedding_network
     self._hidden_layers = hidden_layers
     self._pooler_layer = pooler_layer