Как определить Dense слой для двумерного инпута?

Создаю нейронку (tensorflow 2)

input_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")
segment_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32, name="segment_ids")
bert_inputs = [input_ids, input_mask, segment_ids]

bert_output = BertLayer()(bert_inputs)
dense = tf.keras.layers.Dense(512, activation='softmax')(bert_output)

model = tf.keras.models.Model(inputs=bert_inputs, outputs=dense)
model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
model.summary()

Но на выходе получаю

Incompatible shapes: [32,512] vs. [32,512,512]

Как починить?..