if self.use_proj:
      self.proj_layer = tf_keras.layers.Dense(
          units=self.hidden_size,
          kernel_initializer=self.initializer,
          activation=tf.nn.tanh,