pytorch写一个lstm的model

好的，可以使用以下代码编写一个简单的LSTM模型：

import torch
import torch.nn as nn

class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, batch_first=True):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=batch_first)
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, x):
        # Initialize hidden state with zeros
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).requires_grad_()
        # Initialize cell state with zeros
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).requires_grad_()

        # Initialize hidden state and cell state
        out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach()))

        # Decode the hidden state of the last time step
        out = out[:, -1, :]
        out = self.fc(out)
        return out

其中，`input_size`是输入的特征维度，`hidden_size`是LSTM的隐藏状态维度，`num_layers`是LSTM的层数。需要注意的是，LSTM的输入数据形状应该是`(batch_size, sequence_length, input_size)`，输出数据形状是`(batch_size, 1)`，因此在`forward`函数中需要对最后一个时间步的输出进行处理。