from transformers import DefaultDataCollator
data_collator = DefaultDataCollator(return_tensors="tf")
tf_train_dataset = train_ds.to_tf_dataset(
     columns=["pixel_values", "label"],
     shuffle=True,
     batch_size=batch_size,
     collate_fn=data_collator,
 )
tf_eval_dataset = test_ds.to_tf_dataset(
     columns=["pixel_values", "label"],
     shuffle=True,
     batch_size=batch_size,
     collate_fn=data_collator,
 )