title: Reuse the weight matrix from embedding layer with @tf.function
body: <p>Without using @tf.function, the script work perfectly</p>
  
  <p>I want to use it to speed up training, but it's giving me error where I reuse the weight matrix from the embedding layers. </p>
  
  <p>I think the error is caused by get_weights(), because it converts tensor back to numpy </p>
  
  <p>I tried to use a tf.keras.layers.Dense instead of re-using the weights from embedding, and it worked perfectly.</p>
  
  <pre><code>class Example(tf.keras.Model):
      def __init__(self,):
          super(Example, self).__init__()
          self.embed_dim = embed_dim
          self.vocab_size = vocab_size
          self.embed = tf.keras.layers.Embedding(self.vocab_size, self.embed_dim)
          ...
  
      def call(self, inputs, trianing):
          ...
          embed_matrix = self.embed.get_weights()
  
          # a dense layer
          Vhid = tf.matmul(self.kernel, tf.transpose(embed_matrix[0]))
          pred_w = tf.matmul(pred, Vhid) + self.bias
  </code></pre>
  
  <p>In my train script.
  I did</p>
  
  <pre><code>@tf.function
  def train_step(x, y, training=None):
  
      with tf.GradientTape() as tape:
          pred = model(x, y, training)
          losses = compute_loss(y, pred)
  
      grads = tape.gradient(losses, model.trainable_variables)
  
      optimizer.apply_gradients(zip(grads, model.trainable_variables))
  
      return losses
  
  /home/thomas/projects/tf_convsent/models/.py:195 call  *
     embed_matrix = self.embed.get_weights()  # [vocab_size, 300]
  /home/thomas/.conda/envs/tf2_p37/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:1177 get_weights
     return backend.batch_get_value(params)
  /home/thomas/.conda/envs/tf2_p37/lib/python3.7/site-packages/tensorflow/python/keras/backend.py:3011 batch_get_value
     raise RuntimeError('Cannot get value inside Tensorflow graph function.')
  
  RuntimeError: Cannot get value inside Tensorflow graph function.
  </code></pre>