I've tried it wit `num_classes` doing everything before, but I've decided to add this extra parameter for the reasons above (and maybe I forgot something as well). Also, I think it's easier to explain this to the users, rather than adding some convoluted discussion on how `num_classes` can be used.