LSTM to GRU sampling model issue

Question

I am following this (https://blog.keras.io/a-ten-minute-introduction-to-sequence-to-sequence-learning-in-keras.html) sequence-to-sequence model tutorial. However, when I try to use GRU instead of LSTM in the model. I got the following error.

# GRU
# Define an input sequence and process it.
encoder_inputs = keras.Input(shape=(None, num_encoder_tokens))
encoder = keras.layers.GRU(latent_dim, return_state=True)
encoder_outputs, state_h = encoder(encoder_inputs)

# We discard `encoder_outputs` and only keep the states.
encoder_states = [state_h]

# Set up the decoder, using `encoder_states` as initial state.
decoder_inputs = keras.Input(shape=(None, num_decoder_tokens))

# We set up our decoder to return full output sequences,
# and to return internal states as well. We don't use the
# return states in the training model, but we will use them in inference.
decoder_gru = keras.layers.GRU(latent_dim, return_sequences=True, return_state=True)
decoder_outputs, _= decoder_gru(decoder_inputs, initial_state=encoder_states)
decoder_dense = keras.layers.Dense(num_decoder_tokens, activation="softmax")
decoder_outputs = decoder_dense(decoder_outputs)

# Define the model that will turn
# `encoder_input_data` & `decoder_input_data` into `decoder_target_data`
model = keras.Model([encoder_inputs, decoder_inputs], decoder_outputs)

model.summary()

model.compile(optimizer="rmsprop", loss="categorical_crossentropy", metrics=["accuracy"])


model.fit([train_encoder_input_data, train_decoder_input_data],
train_decoder_target_data,
batch_size=batch_size,
epochs=epochs,
validation_split=0.2,)

model.save("s2s")

model = keras.models.load_model("s2s")


# Define sampling models
# Restore the model and construct the encoder and decoder.
encoder_inputs = model.input[0]  # input_1
encoder_outputs, state_h_enc = model.layers[2].output  # gru_1
encoder_states = [state_h_enc]
encoder_model = keras.Model(encoder_inputs, encoder_states)

decoder_inputs = model.input[1]  # input_2
decoder_state_input_h = keras.Input(shape=(latent_dim,))
decoder_states_inputs = [decoder_state_input_h]
decoder_lstm = model.layers[3]
decoder_outputs, state_h_dec =decoder_lstm(decoder_inputs,initial_state=decoder_states_inputs)
decoder_states = [state_h_dec]
decoder_dense = model.layers[4]
decoder_outputs = decoder_dense(decoder_outputs)
decoder_model = keras.Model([decoder_inputs] + decoder_states_inputs, [decoder_outputs] + decoder_states)

# Reverse-lookup token index to decode sequences back to
# something readable.
reverse_input_char_index = dict((i, char) for char, i in input_token_index.items())
reverse_target_char_index = dict((i, char) for char, i in target_token_index.items())
print(reverse_input_char_index)
print(target_token_index)

def decode_sequence(input_seq):
    # Encode the input as state vectors.
    states_value = encoder_model.predict(input_seq)

    # Generate empty target sequence of length 1.
    target_seq = np.zeros((1, 1, num_decoder_tokens))
    # Populate the first character of target sequence with the start character.
    target_seq[0, 0, target_token_index["\t"]] = 1.0

    # Sampling loop for a batch of sequences
    # (to simplify, here we assume a batch of size 1).
    stop_condition = False
    decoded_sentence = ""
    while not stop_condition:
        output_tokens, h = decoder_model.predict([target_seq] + states_value)

        # Sample a token
        sampled_token_index = np.argmax(output_tokens[0, -1, :])
        sampled_char = reverse_target_char_index[sampled_token_index]
        decoded_sentence += sampled_char

        # Exit condition: either hit max length
        # or find stop character.
        if sampled_char == "\n" or len(decoded_sentence) > max_decoder_seq_length:
            stop_condition = True

        # Update the target sequence (of length 1).
        target_seq = np.zeros((1, 1, num_decoder_tokens))
        target_seq[0, 0, sampled_token_index] = 1.0

        # Update states
        states_value = [h]
    return decoded_sentence



correct = 0
n = 20    
for seq_index in range(20):
    # Take one sequence (part of the training set)
    # for trying out decoding.
    input_seq = train_encoder_input_data[seq_index : seq_index + 1]
    decoded_sentence = decode_sequence(input_seq)
    print("-")
    print("Input sentence:", train_input_texts[seq_index])
    print("Decoded sentence:", decoded_sentence.strip())
    target = train_target_texts[seq_index].strip()
    print("Original sentence:", target)
    if decoded_sentence.strip() == target:
        correct+=1

print((correct / n) * 100)

I got the error as below.

I have modify the code regarding the state as GRU has one state less than LSTM. I am not sure what is the problem here and what should I change and modify. I will really appreciate if someone can tell me what is the problem.