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