Remove the text column because the model does not accept raw text as an input:

tokenized_datasets = tokenized_datasets.remove_columns(["text"])

Rename the label column to labels because the model expects the argument to be named labels:

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")