tokenized_datasets = tokenized_datasets.rename_column("label", "labels")

Set the format of the dataset to return PyTorch tensors instead of lists:

tokenized_datasets.set_format("torch")

Then create a smaller subset of the dataset as previously shown to speed up the fine-tuning:

small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))