def call(self, hidden_states, input_tensor, training=False):
    hidden_states = self.dense(hidden_states)
    hidden_states = self.dropout(hidden_states, training=training)
    hidden_states = self.LayerNorm((hidden_states + input_tensor))
    return hidden_states
def call(self, hidden_states, input_tensor, training=False):
    hidden_states = self.dense(hidden_states)
    hidden_states = self.dropout(hidden_states, training=training)
    hidden_states = self.LayerNorm((hidden_states + input_tensor))
    return hidden_states