如何用torch写一个Bi-LSTM网络

import torch
import torch.nn as nn
class BiLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(BiLSTM, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(hidden_size*2, num_classes)  # *2因为有前向和后向两个方向
    def forward(self, x):
        h0 = torch.zeros(self.num_layers*2, x.size(0), self.hidden_size).to(x.device)  # 初始化隐藏状态
        c0 = torch.zeros(self.num_layers*2, x.size(0), self.hidden_size).to(x.device)  # 初始化单元状态
        # 前向传播
        out, _ = self.lstm(x, (h0, c0))
        # 取得双向LSTM最后一层所有步骤的输出
        out = torch.cat((out[:, -1, :self.hidden_size], out[:, 0, self.hidden_size:]), dim=1)
        # 全连接层
        out = self.fc(out)
        return out

nn.LSTM(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0.0, bidirectional=False, proj_size=0)