ireland-news-headlines-year.../hf_roberta_base_classification/classification_head.py

11 lines
628 B
Python
Raw Normal View History

2021-12-29 13:40:11 +01:00
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