| import torch | |
| import torchvision.models as models | |
| import torch.nn as nn | |
| from huggingface_hub import PyTorchModelHubMixin | |
| class MyResNet18(nn.Module, PyTorchModelHubMixin): | |
| def __init__(self, num_classes): | |
| super().__init__() | |
| self.model = models.resnet18(weights=None) | |
| num_ftrs = self.model.fc.in_features | |
| self.model.fc = nn.Linear(num_ftrs, num_classes) | |
| def forward(self, x): | |
| return self.model(x) |