import torch.nn as nn
from torchvision import models

# Base model
from architectures.Model import Model

class WideResNet(Model):
    
    def get(self):
        # wide_resnet50_2 is a common high-performance choice
        m = models.wide_resnet50_2(weights=models.Wide_ResNet50_2_Weights.DEFAULT)
        m.fc = nn.Linear(m.fc.in_features, self.size)
        return m