14 lines
300 B
Python
14 lines
300 B
Python
|
|
import torch.nn as nn
|
|
from torchvision import models
|
|
|
|
# Base model
|
|
from architectures.Model import Model
|
|
|
|
class GoogleNet(Model):
|
|
|
|
def get(self):
|
|
|
|
m = models.googlenet(weights = models.GoogLeNet_Weights.DEFAULT)
|
|
m.fc = nn.Linear(m.fc.in_features, self.size)
|
|
return m |