title: Tensorflow 2: Getting "WARNING:tensorflow:9 out of the last 9 calls to <function> triggered tf.function retracing. Tracing is expensive"
body: <p>I think this error is coming from a problem with shapes, but I have no idea where. The complete error message suggests to do the following: </p>
  
  <blockquote>
    <p>Also, tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. </p>
  </blockquote>
  
  <p>When I enter this argument in the function decorator, it does work. </p>
  
  <pre><code>@tf.function(experimental_relax_shapes=True)
  </code></pre>
  
  <p>What can the cause be? Here's the full code:</p>
  
  <pre><code>import os
  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
  import tensorflow as tf
  print(f'Tensorflow version {tf.__version__}')
  from tensorflow import keras
  from tensorflow.keras.layers import Dense, Conv1D, GlobalAveragePooling1D, Embedding
  import tensorflow_datasets as tfds
  from tensorflow.keras.models import Model
  
  (train_data, test_data), info = tfds.load('imdb_reviews/subwords8k',
                                            split=[tfds.Split.TRAIN, tfds.Split.TEST],
                                            as_supervised=True, with_info=True)
  
  padded_shapes = ([None], ())
  
  train_dataset = train_data.shuffle(25000).\
      padded_batch(padded_shapes=padded_shapes, batch_size=16)
  test_dataset = test_data.shuffle(25000).\
      padded_batch(padded_shapes=padded_shapes, batch_size=16)
  
  n_words = info.features['text'].encoder.vocab_size
  
  
  class ConvModel(Model):
      def __init__(self):
          super(ConvModel, self).__init__()
          self.embe = Embedding(n_words, output_dim=16)
          self.conv = Conv1D(32, kernel_size=6, activation='elu')
          self.glob = GlobalAveragePooling1D()
          self.dens = Dense(2)
  
      def call(self, x, training=None, mask=None):
          x = self.embe(x)
          x = self.conv(x)
          x = self.glob(x)
          x = self.dens(x)
          return x
  
  
  conv = ConvModel()
  
  conv(next(iter(train_dataset))[0])
  
  loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
  
  train_loss = tf.keras.metrics.Mean()
  test_loss = tf.keras.metrics.Mean()
  
  train_acc = tf.keras.metrics.CategoricalAccuracy()
  test_acc = tf.keras.metrics.CategoricalAccuracy()
  
  optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
  
  
  @tf.function
  def train_step(inputs, labels):
      with tf.GradientTape() as tape:
          logits = conv(inputs, training=True)
          loss = loss_object(labels, logits)
          train_loss(loss)
          train_acc(logits, labels)
  
      gradients = tape.gradient(loss, conv.trainable_variables)
      optimizer.apply_gradients(zip(gradients, conv.trainable_variables))
  
  
  @tf.function
  def test_step(inputs, labels):
      logits = conv(inputs, training=False)
      loss = loss_object(labels, logits)
      test_loss(loss)
      test_acc(logits, labels)
  
  
  def learn():
      train_loss.reset_states()
      test_loss.reset_states()
      train_acc.reset_states()
      test_acc.reset_states()
  
      for text, target in train_dataset:
          train_step(inputs=text, labels=target)
  
      for text, target in test_dataset:
          test_step(inputs=text, labels=target)
  
  
  def main(epochs=2):
      for epoch in tf.range(1, epochs + 1):
          learn()
          template = 'TRAIN LOSS {:&gt;5.3f} TRAIN ACC {:.2f} TEST LOSS {:&gt;5.3f} TEST ACC {:.2f}'
  
          print(template.format(
              train_loss.result(),
              train_acc.result(),
              test_loss.result(),
              test_acc.result()
          ))
  
  if __name__ == '__main__':
      main(epochs=1)
  </code></pre>