def call(self, generator_hidden_states, training=False):
    hidden_states = self.dense(generator_hidden_states)
    hidden_states = get_tf_activation('gelu')(hidden_states)
    hidden_states = self.LayerNorm(hidden_states)
    return hidden_states
def call(self, generator_hidden_states, training=False):
    hidden_states = self.dense(generator_hidden_states)
    hidden_states = get_tf_activation('gelu')(hidden_states)
    hidden_states = self.LayerNorm(hidden_states)
    return hidden_states