Fix get_config() in ReZeroTransformer.

PiperOrigin-RevId: 307297217

Fix get_config() in ReZeroTransformer.
PiperOrigin-RevId: 307297217
295b3b4f · Le Hou · A. Unique TensorFlower · 4e434726 · 295b3b4f
Commit 295b3b4f authored Apr 19, 2020 by Le Hou Committed by A. Unique TensorFlower Apr 19, 2020
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

official/nlp/modeling/layers/rezero_transformer.py official/nlp/modeling/layers/rezero_transformer.py +2 -2

No files found.
--- a/official/nlp/modeling/layers/rezero_transformer.py
+++ b/official/nlp/modeling/layers/rezero_transformer.py
@@ -187,6 +187,8 @@ class ReZeroTransformer(tf.keras.layers.Layer):
            self._dropout_rate,
        "attention_dropout_rate":
            self._attention_dropout_rate,
+        "use_layer_norm":
+            self._use_layer_norm,
        "kernel_initializer":
            tf.keras.initializers.serialize(self._kernel_initializer),
        "bias_initializer":
@@ -201,8 +203,6 @@ class ReZeroTransformer(tf.keras.layers.Layer):
            tf.keras.constraints.serialize(self._kernel_constraint),
        "bias_constraint":
            tf.keras.constraints.serialize(self._bias_constraint),
-        "use_layer_norm":
-            tf.keras.constraints.serialize(self._use_layer_norm)
    }
    base_config = super(ReZeroTransformer, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))