import torch.nn as nn
from torchvision import models

# Base model
from architectures.Model import Model

class ResNet34(Model):
    
    def get(self):
        m = models.resnet34(weights=models.ResNet34_Weights.DEFAULT)

        # freez all layers
        #for param in m.parameters():
        #    param.requires_grad = False

        # unfreez the last two
        #for param in m.layer3.parameters(): param.requires_grad = True
        #for param in m.layer4.parameters(): param.requires_grad = True
        
        m.fc = nn.Linear(m.fc.in_features, self.size)
        return m