编写class LSTM_LM(nn.Module):

好的，下面是一个简单的 `LSTM_LM` 类的实现，用于构建一个基于 LSTM 的语言模型：

import torch
import torch.nn as nn

class LSTM_LM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1, batch_first=True):
        super(LSTM_LM, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.num_layers = num_layers
        self.batch_first = batch_first

        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=batch_first)
        self.linear = nn.Linear(hidden_size, output_size)

    def forward(self, input_seq, hidden_state=None):
        lstm_out, hidden_state = self.lstm(input_seq, hidden_state)
        output = self.linear(lstm_out)

        return output, hidden_state

这个类接受以下参数：

- `input_size`：LSTM 的输入特征数量。
- `hidden_size`：LSTM 的隐藏单元数量。
- `output_size`：模型的输出特征数量。
- `num_layers`：LSTM 的层数，默认为 1。
- `batch_first`：输入张量的第一个维度是否为 batch 的大小，默认为 True。

`LSTM_LM` 类包含一个 LSTM 层和一个线性层，输入是一个张量序列 `input_seq` 和一个 LSTM 隐藏状态 `hidden_state`（默认为 None）。在前向传播中，`input_seq` 会被送入 LSTM 层，得到 LSTM 输出 `lstm_out` 和最后一个时间步的隐藏状态 `hidden_state`。`lstm_out` 会被送入线性层，得到模型的输出 `output`。

希望这个实现能够帮助你！