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