写一个class LSTM(nn.Module)代码

好的，以下是一个简单的 LSTM 模型的代码：

import torch.nn as nn

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(LSTM, self).__init__()
        
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        # 定义 LSTM 层
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        
        # 定义全连接层
        self.fc = nn.Linear(hidden_size, output_size)
        
    def forward(self, x):
        # 初始化隐状态和细胞状态
        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)
        
        # 前向传播 LSTM 层
        out, _ = self.lstm(x, (h0, c0))
        
        # 取出最后一个时间步的输出
        out = out[:, -1, :]
        
        # 进行全连接层计算
        out = self.fc(out)
        
        return out

这个代码实现了一个有一个 LSTM 层和一个全连接层的模型。在初始化函数中，我们定义了 LSTM 层和全连接层。在 forward 函数中，我们首先初始化了隐状态和细胞状态，然后将输入数据传入 LSTM 层中进行前向传播，最后取出最后一个时间步的输出，并将其输入到全连接层中进行计算。