rnn_model.py

HIDDEN_SIZE = 100


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = torch.nn.LSTM(input_size=300, hidden_size=HIDDEN_SIZE)
        self.fc = torch.nn.Linear(HIDDEN_SIZE, 1)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x):
        _, (x, _) = self.lstm(x)
        x = self.fc(x)
        return self.sigmoid(x)