ptrblck:

class Net(nn.Module):
    def __init__(self, input_dim, hidden_dim, layer_dim, bidirectional, output_dim, dropout=0.3):
        super(Net, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.layer_dim = layer_dim
        self.LSTM = nn.GRU(input_dim, hidden_dim, layer_dim, batch_first=True, bidirectional=bidirectional, dropout=dropout)
        self.out = nn.Linear(hidden_dim*2, output_dim, ) if bidirectional else nn.Linear(hidden_dim, output_dim, )

    def forward(self, x):
        # x: (batch, time_step, input_size)
        # LSTM_out: (batch, time_step, output_size)
        # h_n: (n_layers, batch, hidden_size)
        # h_c: (n_layers, batch, hidden_size)
        LSTM_out, h_n = self.LSTM(x, None)
        out = self.out(LSTM_out)
        return out

model = Net(
    input_dim=1,
    hidden_dim=128,
    layer_dim=5,
    output_dim=4,
    bidirectional=True).cuda()


x = torch.randn(128, 800, 1).cuda()
criterion = nn.MSELoss()

out = model(x)
loss = criterion(out, torch.rand_like(out))
loss.backward()

print('done')



Sorry, I forgot to mention, the error will be met after few epochs. I will give you a toy example soon.