16 lines
484 B
Python
16 lines
484 B
Python
|
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
|