Is your feature request related to a problem? Please describe.
Hogler raised this feature request internally, need to support TorchVison models with specified number of classes.