Convert your datasets to the tf.data.Dataset format with [~transformers.TFPreTrainedModel.prepare_tf_dataset]:

tf_train_set = model.prepare_tf_dataset(
     tokenized_squad["train"],
     shuffle=True,
     batch_size=16,
     collate_fn=data_collator,
 )
tf_validation_set = model.prepare_tf_dataset(
     tokenized_squad["test"],
     shuffle=False,
     batch_size=16,
     collate_fn=data_collator,
 )

Configure the model for training with compile: