| class AI(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.embai = nn.Embedding(230, 256) | |
| self.lsai = nn.LSTM(256, 512, batch_first=True, dropout=0.3, num_layers=1) | |
| self.linai = nn.Linear(512, 230) | |
| def forward(self, x): | |
| x = self.embai(x) | |
| if x.dim() == 2: | |
| x = x.unsqueeze(0) | |
| x, _ = self.lsai(x) | |
| x = x[:, -1, :] | |
| x = self.linai(x) | |
| return x |