Spaces:
Build error
Build error
File size: 537 Bytes
edc8afb 295828c edc8afb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
import timm
import torch.nn as nn
from torchvision import models
class ResnetModel(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
model = models.resnet18()
model.fc = nn.Linear(512, 10)
self.model = model
def forward(self, x):
out = self.model(x)
return out
class EffnetModel(nn.Module):
def __init__(self, num_classes=10) -> None:
super().__init__()
model = timm.create_model('efficientnet_b0', num_classes=10)
self.model = model
def forward(self, x):
out = self.model(x)
return out
|