11 lines
628 B
Python
11 lines
628 B
Python
|
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
|