写一个时序预测的class LSTM(nn.Module)代码

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

import torch
import torch.nn as nn

class LSTM(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(LSTM, self).__init__()
        
        self.hidden_dim = hidden_dim
        
        # LSTM layer
        self.lstm = nn.LSTM(input_dim, hidden_dim)
        
        # Output layer
        self.linear = nn.Linear(hidden_dim, output_dim)
        
    def forward(self, input_seq):
        # Initialize hidden state and cell state
        h0 = torch.zeros(1, input_seq.size(1), self.hidden_dim).to(input_seq.device)
        c0 = torch.zeros(1, input_seq.size(1), self.hidden_dim).to(input_seq.device)
        
        # LSTM layer
        lstm_out, _ = self.lstm(input_seq, (h0, c0))
        
        # Only keep the last output of the sequence
        last_out = lstm_out[-1]
        
        # Output layer
        output = self.linear(last_out)
        
        return output

这个模型的输入是一个序列数据，输出是序列最后一个时间步的预测结果。其中：

- `input_dim` 表示输入数据的特征维度；
- `hidden_dim` 表示 LSTM 隐藏状态的维度；
- `output_dim` 表示输出数据的维度。

在 `__init__` 方法中，我们定义了一个 LSTM 层和一个线性输出层。在 `forward` 方法中，我们首先初始化 LSTM 的隐藏状态和细胞状态，然后将输入序列和隐藏状态传入 LSTM 层，得到 LSTM 的输出。我们只保留序列最后一个时间步的输出作为输出层的输入，最终得到预测结果。