Commit 68039be4 authored by Jinoo Baek's avatar Jinoo Baek Committed by A. Unique TensorFlower
Browse files

Change BertPretrainer to keep `mlm_activation` in config.

PiperOrigin-RevId: 434014162
parent 8f815365
...@@ -199,6 +199,7 @@ class BertPretrainerV2(tf.keras.Model): ...@@ -199,6 +199,7 @@ class BertPretrainerV2(tf.keras.Model):
self._config = { self._config = {
'encoder_network': encoder_network, 'encoder_network': encoder_network,
'mlm_initializer': mlm_initializer, 'mlm_initializer': mlm_initializer,
'mlm_activation': mlm_activation,
'classification_heads': classification_heads, 'classification_heads': classification_heads,
'name': name, 'name': name,
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment