title: Exception encountered when calling layer "preprocessing" (type KerasLayer)
body: <p>I try to adapt the BERT tutorial (<a href="https://colab.research.google.com/github/tensorflow/text/blob/master/docs/tutorials/classify_text_with_bert.ipynb" rel="nofollow noreferrer">https://colab.research.google.com/github/tensorflow/text/blob/master/docs/tutorials/classify_text_with_bert.ipynb</a>) where I need to replace the loading of the examples from files to lists and change from binary classification to categorical classification.</p>
  <pre><code>import tensorflow as tf
  import tensorflow_hub as hub
  import tensorflow_text as text
  from official.nlp import optimization
  
  
  target_names = ['cat1', 'cat2', 'cat3']
  x_train = ['Text to categorize 1', 'Text to categorize 2', 'Text to categorize 3', 'Text to categorize 4']
  y_train = [[1, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]]
  x_val = ['Text to categorize 5']
  y_val = [[1, 0, 0]]
  x_test = ['Text to categorize 6']
  y_test = [[0, 1, 0]]
  
  
  AUTOTUNE = tf.data.AUTOTUNE
  batch_size = 32
  seed = 42
  
  x_train = tf.convert_to_tensor(x_train, dtype=tf.string)
  y_train = tf.convert_to_tensor(y_train, dtype=tf.float32)
  x_val = tf.convert_to_tensor(x_val, dtype=tf.string)
  y_val = tf.convert_to_tensor(y_val, dtype=tf.float32)
  x_test = tf.convert_to_tensor(x_test, dtype=tf.string)
  y_test = tf.convert_to_tensor(y_test, dtype=tf.float32)
  
  train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
  train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
  
  val_ds = tf.data.Dataset.from_tensor_slices((x_val, y_val))
  val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
  
  test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
  test_ds = test_ds.cache().prefetch(buffer_size=AUTOTUNE)
  
  del x_train
  del x_val
  del x_test
  del y_train
  del y_val
  del y_test
  
  tf.get_logger().setLevel('ERROR')
  
  bert_model_name = 'bert_en_cased_L-12_H-768_A-12'
  tfhub_handle_encoder = 'https://tfhub.dev/tensorflow/bert_en_cased_L-12_H-768_A-12/3'
  tfhub_handle_preprocess = 'https://tfhub.dev/tensorflow/bert_en_cased_preprocess/3'
  
  bert_preprocess_model = hub.KerasLayer(tfhub_handle_preprocess)
  bert_model = hub.KerasLayer(tfhub_handle_encoder)
  
  
  def build_classifier_model():
      text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name='text')
      preprocessing_layer = hub.KerasLayer(tfhub_handle_preprocess, name='preprocessing')
      encoder_inputs = preprocessing_layer(text_input)
      encoder = hub.KerasLayer(tfhub_handle_encoder, trainable=True, name='BERT_encoder')
      outputs = encoder(encoder_inputs)
      net = outputs['pooled_output']
      net = tf.keras.layers.Dense(64, activation='relu')(net)
      net = tf.keras.layers.Dropout(0.2)(net)
      net = tf.keras.layers.Dense(32, activation='relu')(net)
      net = tf.keras.layers.Dropout(0.2)(net)
      net = tf.keras.layers.Dense(len(target_names), activation='softmax', name='classifier')(net)
      return tf.keras.Model(text_input, net)
  
  
  classifier_model = build_classifier_model()
  
  loss = tf.keras.losses.CategoricalCrossentropy()
  metrics = tf.metrics.CategoricalAccuracy()
  
  epochs = 5
  steps_per_epoch = tf.data.experimental.cardinality(train_ds).numpy()
  num_train_steps = steps_per_epoch * epochs
  num_warmup_steps = int(0.1*num_train_steps)
  
  init_lr = 3e-5
  optimizer = optimization.create_optimizer(init_lr=init_lr,
                                            num_train_steps=num_train_steps,
                                            num_warmup_steps=num_warmup_steps,
                                            optimizer_type='adamw')
  
  classifier_model.compile(optimizer=optimizer,
                           loss=loss,
                           metrics=metrics)
  
  history = classifier_model.fit(x=train_ds,
                                 validation_data=val_ds,
                                 epochs=epochs)
  
  classifier_model.save('bert', include_optimizer=False)
  </code></pre>
  <p>My code seems to work until the 'classifier_model.fit', then I get an error that I don't understand :(</p>
  <pre><code>Epoch 1/5
  Traceback (most recent call last):
    File &quot;/home/.../venv/lib64/python3.7/site-packages/keras/utils/traceback_utils.py&quot;, line 67, in error_handler
      raise e.with_traceback(filtered_tb) from None
    File &quot;/tmp/__autograph_generated_file2izv5eiz.py&quot;, line 15, in tf__train_function
      retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
    File &quot;/tmp/__autograph_generated_fileffjts0m6.py&quot;, line 74, in tf__call
      ag__.if_stmt(ag__.not_(ag__.ld(self)._has_training_argument), if_body_3, else_body_3, get_state_3, set_state_3, ('result', 'training'), 1)
    File &quot;/tmp/__autograph_generated_fileffjts0m6.py&quot;, line 72, in else_body_3
      result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)
    File &quot;/tmp/__autograph_generated_fileffjts0m6.py&quot;, line 72, in &lt;lambda&gt;
      result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)
  ValueError: in user code:
  
      File &quot;/home/.../venv/lib64/python3.7/site-packages/keras/engine/training.py&quot;, line 1051, in train_function  *
          return step_function(self, iterator)
      File &quot;/home/.../venv/lib64/python3.7/site-packages/keras/engine/training.py&quot;, line 1040, in step_function  **
          outputs = model.distribute_strategy.run(run_step, args=(data,))
      File &quot;/home/.../venv/lib64/python3.7/site-packages/keras/engine/training.py&quot;, line 1030, in run_step  **
          outputs = model.train_step(data)
      File &quot;/home/.../venv/lib64/python3.7/site-packages/keras/engine/training.py&quot;, line 889, in train_step
          y_pred = self(x, training=True)
      File &quot;/home/.../venv/lib64/python3.7/site-packages/keras/utils/traceback_utils.py&quot;, line 67, in error_handler
          raise e.with_traceback(filtered_tb) from None
      File &quot;/tmp/__autograph_generated_fileffjts0m6.py&quot;, line 74, in tf__call
          ag__.if_stmt(ag__.not_(ag__.ld(self)._has_training_argument), if_body_3, else_body_3, get_state_3, set_state_3, ('result', 'training'), 1)
      File &quot;/tmp/__autograph_generated_fileffjts0m6.py&quot;, line 72, in else_body_3
          result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)
      File &quot;/tmp/__autograph_generated_fileffjts0m6.py&quot;, line 72, in &lt;lambda&gt;
          result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)
  
      ValueError: Exception encountered when calling layer &quot;preprocessing&quot; (type KerasLayer).
      
      in user code:
      
          File &quot;/home/.../venv/lib64/python3.7/site-packages/tensorflow_hub/keras_layer.py&quot;, line 237, in call  *
              result = smart_cond.smart_cond(training,
      
          ValueError: Could not find matching concrete function to call loaded from the SavedModel. Got:
            Positional arguments (3 total):
              * &lt;tf.Tensor 'inputs:0' shape=() dtype=string&gt;
              * False
              * None
            Keyword arguments: {}
          
           Expected these arguments to match one of the following 4 option(s):
          
          Option 1:
            Positional arguments (3 total):
              * TensorSpec(shape=(None,), dtype=tf.string, name='inputs')
              * True
              * None
            Keyword arguments: {}
          
          Option 2:
            Positional arguments (3 total):
              * TensorSpec(shape=(None,), dtype=tf.string, name='sentences')
              * True
              * None
            Keyword arguments: {}
          
          Option 3:
            Positional arguments (3 total):
              * TensorSpec(shape=(None,), dtype=tf.string, name='inputs')
              * False
              * None
            Keyword arguments: {}
          
          Option 4:
            Positional arguments (3 total):
              * TensorSpec(shape=(None,), dtype=tf.string, name='sentences')
              * False
              * None
            Keyword arguments: {}
      
      
      Call arguments received by layer &quot;preprocessing&quot; (type KerasLayer):
        • inputs=tf.Tensor(shape=(), dtype=string)
        • training=True
  
  python-BaseException
  
  Process finished with exit code 1
  </code></pre>