Minor bug fixes

PiperOrigin-RevId: 335944731

Minor bug fixes
PiperOrigin-RevId: 335944731
d4557be2 · Le Hou · A. Unique TensorFlower · 32d5a722 · d4557be2
Commit d4557be2 authored Oct 07, 2020 by Le Hou Committed by A. Unique TensorFlower Oct 07, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

official/nlp/modeling/layers/transformer_scaffold.py official/nlp/modeling/layers/transformer_scaffold.py +2 -0

No files found.
--- a/official/nlp/modeling/layers/transformer_scaffold.py
+++ b/official/nlp/modeling/layers/transformer_scaffold.py
@@ -224,6 +224,8 @@ class TransformerScaffold(tf.keras.layers.Layer):
            self._dropout_rate,
        "attention_dropout_rate":
            self._attention_dropout_rate,
+        "norm_first":
+            self._norm_first,
        "kernel_initializer":
            tf.keras.initializers.serialize(self._kernel_initializer),
        "bias_initializer":