| import torch | |
| import torch.nn as nn | |
| class SimpleRNN(nn.Module): | |
| def __init__(self, input_size, hidden_size, output_size): | |
| super(SimpleRNN, self).__init__() | |
| self.input_size = input_size | |
| self.hidden_size = hidden_size | |
| self.rnn = nn.RNN(input_size, hidden_size, batch_first=True) | |
| self.fc = nn.Linear(hidden_size, output_size) | |
| def forward(self, x, hidden): | |
| x = torch.nn.functional.one_hot(x, num_classes=self.input_size).float() | |
| out, hidden = self.rnn(x.unsqueeze(0), hidden) | |
| out = self.fc(out[:, -1, :]) | |
| return out, hidden | |