List index out of range when saving a fine tuned bert model

Question

Model is created using the following function definition

def create_model(max_length = 256):
  bert_model = TFBertModel.from_pretrained('bert-base-uncased')
  for layer in bert_model.layers:
    layer.trainable = False
  input_ids = tf.keras.Input(shape = (max_length, ), dtype = tf.int32, name = 'input_ids')
  attention_masks = tf.keras.Input(shape = (max_length, ), dtype = tf.int32, name = 'attention_masks')
  x = bert_model.bert([input_ids, attention_masks])
  x = x.pooler_output
  x = tf.keras.layers.Dropout(0.2)(x)
  x = tf.keras.layers.Dense(256, activation = 'relu')(x)
  x = tf.keras.layers.Dropout(0.2)(x)
  x = tf.keras.layers.Dense(33)(x)
  out = tf.keras.layers.Activation('sigmoid')(x)
  model = tf.keras.Model(inputs = [input_ids, attention_masks], outputs = out)
  model.compile(optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5),
                loss = tf.keras.losses.BinaryCrossentropy(),
                metrics = tf.metrics.BinaryAccuracy())
  return model

On trying to save the model using tf.keras.models.save_model, I run into the following error:
IndexError: Exception encountered when calling layer 'bert' (type TFBertMainLayer). list index out of range