Hi everyone, I finally found a solution.

Do not use this line

# Keras will expect a tuple when dealing with labels
train_dataset = train_dataset.map(lambda x, y: (x, (y['start_positions'], y['end_positions'])))

in graph mode, the model returns the prediction as “TFQuestionAnsweringModelOutput” and this model is inherited from a dictionary! So we need to use the “y” value as a dictionary too.

Replace every “start_positions” with “start_logits” and every “end_positions” with “end_logits”.

The reason is when you calculate loss
TFQuestionAnsweringModelOutput has  “start_logits” and "end_logits"
but
y_true has “start_positions” and “end_positions”.
After the loss calculation, tensorflow adds  “start_logits” and “end_logits”  keys to the y_true and throws an error about length of the y_pred and y_true is different. When you replace start_positions with start_logits etc. problem solved!
NOTE: When you apply step 2, your training code will not work on PyTorch training  You need to set the dictionary keys as   “start_positions” and “end_positions”.