Convert Keras .h5 model to TFLite .tflite with TensorFlow v2.2

I am trying to convert a network I defined using Keras to tflite. The network is as follows:

model = tf.keras.Sequential([
        # Embedding
        tf.keras.layers.Embedding(vocab_size, embedding_dim, batch_input_shape=[BATCH_SIZE, None]),
        # GRU unit
        tf.keras.layers.GRU(rnn_units, return_sequences=True, stateful=True, recurrent_initializer='glorot_uniform'),
        # Fully connected layer
        tf.keras.layers.Dense(vocab_size)
        ])

However, when I try to export to .tflite, it looks like something goes wrong due to the presence of the GRU layer.

# Save trained model in .h5 format
keras_file = 'inference.h5'
tf.keras.models.save_model(model, keras_file)

# Load .h5 model with custom loss function
model = load_model('inference.h5', custom_objects={'loss': loss})

# Converting a tf.Keras model to a TensorFlow Lite model.
converter    = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

with error:

ValueError: Input 0 of node sequential_8/gru_2/AssignVariableOp was passed float from sequential_8/gru_2/68029:0 incompatible with expected resource.

Any solution to solve this issue?