import torch.nn as nn
from torchvision import models
from architectures.Model import Model

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

        # load pretrained
        m = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
        # will modify only the final layers
        num_ftrs = m.classifier.in_features
        m.classifier = nn.Linear(num_ftrs, self.size)
        
        return m.to(self.device)