How can I load an exported model (or .pth file) for use on a different dataset (with a different number of classes)?

I trained a model and am now trying to use it on a different, but similar, dataset. However, I am getting an error because there is a different amount of classes on the new dataset. How can I load it? (as I would any other pre-trained model)


You don’t specify which type of model you have, maybe this thread below will help. Otherwise look for other threads or the docs for “custom head”, you have to basically remove the classifier head again and replace it with a new head with the desired num of classes.

1 Like

Yup, that did the trick. Thanks!