import torch.nn as nn
from torchvision import models

# Base model
from architectures.Model import Model

class EfficientNet(Model):
    
    def get(self):

        m = models.efficientnet_b1(weights=models.EfficientNet_B1_Weights.DEFAULT)

        # Unfreeze the last block 
       # for param in m.features[-1].parameters(): param.requires_grad = True

        # Standard classifier fix
        m.classifier[1] = nn.Linear(m.classifier[1].in_features, self.size)
        return m