Solved by nielsr in post #2 
              
              
                Hi, 
I see you are first working with a HuggingFace Dataset (that is returned by the load_dataset function), and that you are then converting it to a PyTorch Dataset. 
Actually, the latter is not required. Also, you can tokenize your training and test splits in one go: 
from transformers import Bert…