Yes sorry, I coiped the wrong line but the output I posted was from dataset_valid.

And also, it seems for  predictions , you have used  batch_size=5000  which lead to 5000 predictions but you are using labels of entire dataset which is 10000.

So, If I’ve understood well, random_split splits the input data but not the labels, right?
In this case, what should i do?
Sorry to bother you with maybe dumb questions, I’m new at PyTorch.