I was trying to train a bert model to solve a muti-classification problem:

• I got this error while run the code below

Arguments target and output must have the same shape. Received: target.shape=(None, 512), output.shape=(None, 3)

import tensorflow as tf

epochs = 4

train_dataloader = train_dataset.shuffle(buffer_size=10000).batch(batch_size)
validation_dataloader = val_dataset.batch(batch_size)

# start training 
history = model.fit(
    train_dataloader,  # train_data
    validation_data=validation_dataloader,  # validation_data
    epochs=epochs,  
    verbose=1 
)
# save the model
model.save("bert_model.h5")

• This is a test

for batch in train_dataloader.take(1):
    input_ids, attention_masks, labels = batch
    print("Batch input_ids shape:", input_ids.shape) 
    print("Batch attention_masks shape:", attention_masks.shape) 
    print("Batch labels shape:", labels.shape)  

# I got this output
Batch input_ids shape: (16, 512)
Batch attention_masks shape: (16, 512)
Batch labels shape: (16,)

I already check the tensor shape.
Hope for answer!