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

I was trying to train a BERT model to solve a multi-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 checked the tensor shape.

Your labels have a shape of (16,), while your model's output has a shape of (None,3).

Probably the issue is that your labels are not one-hot encoded. They should have the same second dimension as your output layer:

from tensorflow.keras.utils import to_categorical
num_classes = 3
labels = to_categorical(labels, num_classes=num_classes)
print(labels.shape)