import torch.nn as nn
from torchvision import models

# Base model
from architectures.Model import Model

class ShuffleNet(Model):
    
    def get(self):
        m = models.shufflenet_v2_x1_0(weights=models.ShuffleNet_V2_X1_0_Weights.DEFAULT)
        
        num_ftrs = m.fc.in_features
        m.fc = nn.Linear(num_ftrs, self.size)
        return m