写一个torch版本的LSTM模型代码

在 PyTorch 中使用 LSTM 的代码如下所示：

import torch
import torch.nn as nn

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(LSTM, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, num_classes)
    
    def forward(self, x):
        # Set initial hidden and cell states 
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device) 
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
        
        # Forward propagate LSTM
        out, _ = self.lstm(x, (h0, c0))  # out: tensor of shape (batch_size, seq_length, hidden_size)
        
        # Decode the hidden state of the last time step
        out = self.fc(out[:, -1, :])
        return out

这段代码定义了一个名为 `LSTM` 的 PyTorch 模型，该模型使用 LSTM 层来处理输入序列，并使用一个全连接层来将最后时间步的隐藏状态解码为输出。

参数说明：

- `input_size`：每个时间步输入的特征维度
- `hidden_size`：LSTM 层的隐藏状态的维度
- `num_layers`：LSTM 层的数量
- `num_classes`：输出的类别数量

在调用 `forward` 方法时，需要传入输入张量 `x`，该张量的形状为 `(batch_size, seq_length, input_size)`。模型会处理输入序列，并返回输出张量，形状为 `(batch_size, num_classes)`。

希望这些