帮我用pytorch写一个Bi-GRU

好的，以下是一个简单的Bi-GRU模型的代码：

import torch
import torch.nn as nn

class BiGRU(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(BiGRU, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.output_size = output_size
        
        self.gru = nn.GRU(input_size, hidden_size, num_layers=num_layers, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(hidden_size * 2, output_size)
    
    def forward(self, x):
        h0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
        out, _ = self.gru(x, h0)
        out = self.fc(out[:, -1, :])
        return out

这个模型有四个参数：

- `input_size`：输入特征的维度。
- `hidden_size`：隐藏状态的维度。
- `num_layers`：GRU层数。
- `output_size`：输出的维度。

该模型使用nn.GRU实现双向GRU，然后通过一个全连接层将最后一个时间步的隐藏状态映射到输出空间。