import torch.nn as nn class CharModel(nn.Module): def __init__(self, n_vocab): super().__init__() self.lstm = nn.LSTM(input_size=1, hidden_size=256, num_layers=1, batch_first=True) self.dropout = nn.Dropout(0.2) self.linear = nn.Linear(256, n_vocab) def forward(self, x): x, _ = self.lstm(x) # take only the last output x = x[:, -1, :] # produce output x = self.linear(self.dropout(x)) return x