How to load a pretrained custom model using `from_pretrained`

Beginners
2

Yes you can inherit from PreTrainedModel to inherit methods like from_pretrained, save_pretrained and push_to_hub.

Alternatively, you can leverage the PyTorchModelHubMixin class available in the huggingface_hub library. This allows you to get the same functionality:

from torch import nn
from huggingface_hub import PyTorchModelHubMixin

class CustomModel(nn.Module, PyTorchModelHubMixin):
          ...
2 Likes