Internal change

PiperOrigin-RevId: 305530702

Internal change
PiperOrigin-RevId: 305530702
91495d13 · Chen Chen · A. Unique TensorFlower · f41682df · 91495d13
Commit 91495d13 authored Apr 08, 2020 by Chen Chen Committed by A. Unique TensorFlower Apr 08, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

official/nlp/bert/bert_models.py official/nlp/bert/bert_models.py +2 -1

No files found.
--- a/official/nlp/bert/bert_models.py
+++ b/official/nlp/bert/bert_models.py
@@ -121,7 +121,8 @@ def get_transformer_encoder(bert_config,
        attention_dropout_rate=bert_config.attention_probs_dropout_prob,
    )
    kwargs = dict(embedding_cfg=embedding_cfg, hidden_cfg=hidden_cfg,
-                  num_hidden_instances=bert_config.num_hidden_layers,)
+                  num_hidden_instances=bert_config.num_hidden_layers,
+                  num_output_classes=bert_config.hidden_size)

    # Relies on gin configuration to define the Transformer encoder arguments.
    return transformer_encoder_cls(**kwargs)