BERT Text Classification Model gives error

Question

I am fine-tuning a BERT Model for text classification with Tensorflow. Here is my code for building the model:

# Building the model
def create_model():
  input_word_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                      name="input_word_ids"),
  input_mask = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                  name="input_mask"),
  input_type_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                  name="input_type_ids")
 
  pooled_output, sequence_output = bert_layer([input_word_ids, input_mask, input_type_ids])

  drop = tf.keras.layers.Dropout(0.4)(pooled_output)
  output = tf.keras.layers.Dense(1, activation="sigmoid", name="output")(drop)
  
  model = tf.keras.Model(
      inputs={
          'input_word_ids': input_word_ids,
          'input_mask': input_mask,
          'input_type_ids': input_type_ids
          },
          outputs=output)
  return model

I used this BERT Layer:

bert_layer = hub.KerasLayer("https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/4", trainable=True)

When creating the model: model = create_model(), it gives the following error:

ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (3 total):
     [(<tf.Tensor 'inputs:0' shape=(None, 512) dtype=int32>,), (<tf.Tensor 'inputs_1:0' shape= 
(None, 512) dtype=int32>,), <tf.Tensor 'inputs_2:0' shape=(None, 512) dtype=int32>]
     False
     None
  Keyword arguments: {}

Expected these arguments to match one of the following 4 option(s):

Option 1:
  Positional arguments (3 total):
     {'input_word_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='input_word_ids'), 
'input_type_ids': TensorSpec(shape=(None, None), dtype=tf.int32, name='input_type_ids'), 
'input_mask': TensorSpec(shape=(None, None), dtype=tf.int32, name='input_mask')}
     False
     None
  Keyword arguments: {}

I would appreciate it if you have any solution in mind.