not sure how many labels you have, but its possible that you have several classes that are in the test set and are not in the train set. alternatively, the number of samples you pass to the train/val/test are more you declare. it would be easier if you pointed to the line where the error is..