You need to use DatasetType.Fix here, which is your training set with the validation transforms and shuffle=False, drop_last=False. Your issue is that since we drop the last incomplete batch (here of size 10) you’re missing 10 predictions. The other issue you don’t see is that your predictions are shuffled and not in the same order as your training set.