使用Pytorch的基类“nn.Linear”实现简单的单层RNN很困难

import torch
import torch.nn as nn

# 定义单层RNN模型
class SimpleRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleRNN, self).__init__()
        self.hidden_size = hidden_size

        self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        # 初始化隐藏状态
        h0 = torch.zeros(1, x.size(0), self.hidden_size).to(x.device)

        # 前向传播
        out, _ = self.rnn(x, h0)
        out = self.fc(out[:, -1, :])  # 取最后一个时间步的输出

        return out

# 创建模型实例
input_size = 10
hidden_size = 20
output_size = 5
model = SimpleRNN(input_size, hidden_size, output_size)

# 打印模型结构
print(model)

# 输入数据
batch_size = 3
seq_length = 4
input_data = torch.randn(batch_size, seq_length, input_size)

# 前向传播
output = model(input_data)
print(output)