At this point, you may need to restart your notebook or execute the following code to free some memory:
py
del model
del trainer
torch.cuda.empty_cache()
Next, manually postprocess tokenized_dataset to prepare it for training.