title: Keras loadmodel for custom model with custom layers - Transformer documentation example
body: <p>I am running the following example:</p>
  <p><a href="https://keras.io/examples/nlp/text_classification_with_transformer/" rel="nofollow noreferrer">https://keras.io/examples/nlp/text_classification_with_transformer/</a></p>
  <p>I have created and trained a model as described, and it works nicely:</p>
  <pre><code>inputs = layers.Input(shape=(maxlen,))
  embedding_layer = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)
  x = embedding_layer(inputs)
  transformer_block = TransformerBlock(embed_dim, num_heads, ff_dim)
  x = transformer_block(x,training=True)
  x = layers.GlobalAveragePooling1D()(x)
  x = layers.Dropout(0.1)(x)
  x = layers.Dense(20, activation=&quot;relu&quot;)(x)
  x = layers.Dropout(0.1)(x)
  outputs = layers.Dense(2, activation=&quot;softmax&quot;)(x)
  
  model = keras.Model(inputs=inputs, outputs=outputs)
  
  
  &quot;&quot;&quot;
  ## Train and Evaluate
  &quot;&quot;&quot;
  
  model.compile(&quot;adam&quot;, &quot;sparse_categorical_crossentropy&quot;, metrics=[&quot;accuracy&quot;])
  history = model.fit(
      x_train, y_train, batch_size=1024, epochs=1, validation_data=(x_val, y_val)
  )
  
  model.save('SPAM.h5')
  </code></pre>
  <p>How to save and load such custom model correctly in Keras?</p>
  <p>I have tried</p>
  <pre><code> best_model=tf.keras.models.load_model('SPAM.h5')
   ValueError: Unknown layer: TokenAndPositionEmbedding
  </code></pre>
  <p>but the model seems to miss the custom Layers. However the following also does not work</p>
  <pre><code>best_model=tf.keras.models.load_model('SPAM.h5',custom_objects={&quot;TokenAndPositionEmbedding&quot;: TokenAndPositionEmbedding()})
   
  TypeError: __init__() missing 3 required positional arguments:
   'maxlen', 'vocab_size', and 'embed_dim'
  </code></pre>
  <p>also passing the class does not solve.</p>
  <pre><code>best_model=tf.keras.models.load_model('SPAM.h5',
   custom_objects={&quot;TokenAndPositionEmbedding&quot;: TokenAndPositionEmbedding})
   TypeError: __init__() got an unexpected keyword argument 'name'
  
  
  
   best_model=tf.keras.models.load_model('SPAM.h5',
  {&quot;TokenAndPositionEmbedding&quot;:
  TokenAndPositionEmbedding,'TransformerBlock':TransformerBlock,
  'MultiHeadSelfAttention':MultiHeadSelfAttention})
  </code></pre>