import torch

class YearClassificationHead(torch.nn.Module):
    def __init__(self, in_dim, MIN_YEAR, MAX_YEAR):                                                                                                                                                                                                                                                                                                       
        super(YearClassificationHead, self).__init__()
        self.linear = torch.nn.Linear(in_dim, MAX_YEAR - MIN_YEAR + 1)
    def forward(self, x): 
        x = x.mean(1)
        x = self.linear(x)
        return x