changing pooling output

Author: mbrhd

keras_nlp/models/bert/bert_backbone.py

@@ -162,13 +162,12 @@ def __init__(
         # Construct the two BERT outputs. The pooled output is a dense layer on
         # top of the [CLS] token.
         sequence_output = x
-        x = keras.layers.Dense(
+        pooled_output = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=bert_kernel_initializer(),
             activation="tanh",
             name="pooled_dense",
-        )(x)
-        pooled_output = x[:, cls_token_index, :]
+        )(x[:, cls_token_index, :])
 
         # Instantiate using Functional API Model constructor
         super().__init__(

keras_nlp/models/electra/electra_backbone.py

@@ -163,13 +163,15 @@ def __init__(
             )(x, padding_mask=padding_mask)
 
         sequence_output = x
-        x = keras.layers.Dense(
+        # Construct the two ELECTRA outputs. The pooled output is a dense layer on
+        # top of the [CLS] token.
+        sequence_output = x
+        pooled_output = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=electra_kernel_initializer(),
             activation="tanh",
             name="pooled_dense",
-        )(x)
-        pooled_output = x[:, cls_token_index, :]
+        )(x[:, cls_token_index, :])
 
         # Instantiate using Functional API Model constructor
         super().__init__(