pytorch_lstm_example.py

from torch import nn
from torch.nn.functional import softmax
class CustomModel(nn.Module):
    def __init__(self, vocab_size=50,
                 embedding_dim=16,
                 hidden_size=8):
        super().__init__()
        self.encoder = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_size)
        self.linear = nn.Linear(hidden_size, 1)
	self.activation = nn.Sigmoid()

    def forward(self, x):
        output = self.encoder(x)
        output, _ = self.lstm(output)
        output = output[-1] # Keep last output only
        output = self.linear(output)
	output = self.activation(output)
        return output