"""
Builds Pytorch model
"""
import torch
import torchvision.models
from torch import nn


class ResNet101(nn.Module):
    """
    ResNet101 model specified for the binary problem. The according transforms were taken from pytorch.org.
    """

    def __init__(self):
        super().__init__()
        self.weights = torchvision.models.ResNet101_Weights.DEFAULT
        self.transforms = self.weights.transforms
        self.resnet = torchvision.models.resnet101(weights=self.weights)

        for param in self.resnet.parameters():
            param.requires_grad = False

        self.resnet.fc = nn.Linear(in_features=2048, out_features=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.resnet(x)
        return x