title: Transformer model not able to be saved
body: <p>I'm trying to follow this tutrial <a href="https://colab.research.google.com/github/tensorflow/examples/blob/master/community/en/transformer_chatbot.ipynb" rel="nofollow noreferrer">https://colab.research.google.com/github/tensorflow/examples/blob/master/community/en/transformer_chatbot.ipynb</a>, However, when I tried to save the model in order to load it again without training I got an error mentioned here <a href="https://stackoverflow.com/questions/58678836/notimplementederror-layers-with-arguments-in-init-must-override-get-conf">NotImplementedError: Layers with arguments in `__init__` must override `get_config`</a>
  I understood from the answer that I need to make the encoder and decoder as classes and customise it(instead of leaving it as functions like the colab tutrial) so I went back to tensor flow documentation of this model here: <a href="https://www.tensorflow.org/tutorials/text/transformer#encoder_layer" rel="nofollow noreferrer">https://www.tensorflow.org/tutorials/text/transformer#encoder_layer</a> and tried to edit in it. I made the encoder layer as:</p>
  
  <pre><code>class EncoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads,  rate=0.1,**kwargs,):
      #super(EncoderLayer, self).__init__()
      super().__init__(**kwargs)
      self.mha = MultiHeadAttention(d_model, num_heads)
      self.ffn = point_wise_feed_forward_network(d_model, dff)
  
      self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
      self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
  
      self.dropout1 = tf.keras.layers.Dropout(rate)
      self.dropout2 = tf.keras.layers.Dropout(rate)
    def get_config(self):
  
          config = super().get_config().copy()
          config.update({
              #'vocab_size': self.vocab_size,
              #'num_layers': self.num_layers,
              #'units': self.units,
              'd_model': self.d_model,
              'num_heads': self.num_heads,
              'dropout': self.dropout,
          })
          return config
  
    def call(self, x, training, mask):
  
      attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)
      attn_output = self.dropout1(attn_output, training=training)
      out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)
  
      ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)
      ffn_output = self.dropout2(ffn_output, training=training)
      out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)
  
      return out2
  </code></pre>
  
  <p>and same for the decoder layer class. Then the same encoder in the documentation of tf</p>
  
  <pre><code>class Encoder(tf.keras.layers.Layer):
    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,
                 maximum_position_encoding, rate=0.1):
      super(Encoder, self).__init__()
  
      self.d_model = d_model
      self.num_layers = num_layers
  
      self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)
      self.pos_encoding = positional_encoding(maximum_position_encoding, 
                                              self.d_model)
  
  
      self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) 
                         for _ in range(num_layers)]
  
      self.dropout = tf.keras.layers.Dropout(rate)
  
    def call(self, x, training, mask):
  
      seq_len = tf.shape(x)[1]
  
      # adding embedding and position encoding.
      x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
      x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
      x += self.pos_encoding[:, :seq_len, :]
  
      x = self.dropout(x, training=training)
  
      for i in range(self.num_layers):
        x = self.enc_layers[i](x, training, mask)
  
      return x  # (batch_size, input_seq_len, d_model)
  </code></pre>
  
  <p>the function of the model as:</p>
  
  <pre><code>def transformer(vocab_size,
                  num_layers,
                  units,
                  d_model,
                  num_heads,
                  dropout,
                  name="transformer"):
    inputs = tf.keras.Input(shape=(None,), name="inputs")
    dec_inputs = tf.keras.Input(shape=(None,), name="dec_inputs")
  
    enc_padding_mask = tf.keras.layers.Lambda(
        create_padding_mask, output_shape=(1, 1, None),
        name='enc_padding_mask')(inputs)
    # mask the future tokens for decoder inputs at the 1st attention block
    look_ahead_mask = tf.keras.layers.Lambda(
        create_look_ahead_mask,
        output_shape=(1, None, None),
        name='look_ahead_mask')(dec_inputs)
    # mask the encoder outputs for the 2nd attention block
    dec_padding_mask = tf.keras.layers.Lambda(
        create_padding_mask, output_shape=(1, 1, None),
        name='dec_padding_mask')(inputs)
  
    enc_outputs = Encoder(
        num_layers=num_layers, d_model=d_model, num_heads=num_heads, 
                           input_vocab_size=vocab_size,
  
  
    )(inputs=[inputs, enc_padding_mask])
  
    dec_outputs = Decoder(
        num_layers=num_layers, d_model=d_model, num_heads=num_heads, 
                            target_vocab_size=vocab_size,
  
  
    )(inputs=[dec_inputs, enc_outputs, look_ahead_mask, dec_padding_mask])
  
    outputs = tf.keras.layers.Dense(units=vocab_size, name="outputs")(dec_outputs)
  
    return tf.keras.Model(inputs=[inputs, dec_inputs], outputs=outputs, name=name)
  </code></pre>
  
  <p>and calling the model:</p>
  
  <pre><code>#the model itself with its paramters:
  # Hyper-parameters
  NUM_LAYERS = 3
  D_MODEL = 256
  #D_MODEL=tf.cast(D_MODEL, tf.float32)
  
  NUM_HEADS = 8
  UNITS = 512
  DROPOUT = 0.1
  model = transformer(
      vocab_size=VOCAB_SIZE,
      num_layers=NUM_LAYERS,
      units=UNITS,
      d_model=D_MODEL,
      num_heads=NUM_HEADS,
      dropout=DROPOUT)
  </code></pre>
  
  <p>However, I got that error:
  <code>TypeError: __init__() missing 2 required positional arguments: 'dff' and 'maximum_position_encoding'</code>
  I am really confused and I don't understand what dff and maximum position encoding mean in the documentation and when I removed them from the encoder and decoder classes, I got anther error as positional_encoding function takes maximum position as input and also dff is passed as input inside the class. I am not so sure what I should do as I am not sure whether I am following the right steps or not</p>