Customizable classification

I’m not sure the number of classes affects the model definition itself beyond specifying num_classes when declaring the model: torchvision.models.resnet — Torchvision 0.12 documentation.

If you are referring to splitting the dataset for evaluation, you might need to write a little custom code to do the stratified sampling e.g., How to do a stratified split - PyTorch Forums