Kernel died when fine-tune a pretrained model using keras

Beginners
1

I am following the hugging face tutorial to fine-tune a BERT model for sequence classification on the CoLA (Corpus of Linguistic Acceptability) dataset using TensorFlow, the kernel died when run the model.fit(). Below is my code:

from transformers import AutoTokenizer
from transformers import TFAutoModelForSequenceClassification
from tensorflow.keras.optimizers import Adam
from datasets import load_dataset

dataset = load_dataset("glue", "cola")
dataset = dataset["train"]
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

def tokenize_dataset(data):
    return tokenizer(data["sentence"])


dataset = dataset.map(tokenize_dataset)
model = TFAutoModelForSequenceClassification.from_pretrained("bert-base-cased")

tf_dataset = model.prepare_tf_dataset(dataset, batch_size=2, shuffle=True, tokenizer=tokenizer)
model.compile(optimizer=Adam(3e-5))  # No loss argument!

model.fit(tf_dataset)

I tried to reduce the batch size, didn’t resolve the issue.
And the code works fine if I use the dataset as array instead of tf.dataset.

Does anyone have some thoughts on this?
Thanks