I’m not sure how you are using limited_set. If you could like to only store certain classes, you should add some kind of condition to select the corresponding data samples and labels.
That being said, try to add an assert statement in the training loop and check for the invalid values, i.e. (labels < 0).any() and (labels >= nb_limited_classes).any(). This should raise the exception when the invalid batch is being processed.
Then check the values of labels and make sure they only contain class indices in the required range.