LSTM Архитектура

Question

Я пытаюсь понять, похож ли мой двухслойный LSTM на изображение.Кто-то может мне помочь?ps я очистил средний слой пула.

входной слой (зеленый), скрытый слой LSTM (синий), линейный слой (желтый)

https://ibb.co/1qN6q1m

class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(RNN, self).__init__()

        self.hidden_size = hidden_size
        self.num_layer=2
        self.i2h = nn.LSTM(input_size, hidden_size, self.num_layer, dropout=0.5)

        self.i2o = nn.Linear(hidden_size, output_size)
        self.softmax = nn.LogSoftmax(dim=1)
        self.hidden = self.initHidden()

    def forward(self, input):
        lstmout, self.hidden= self.i2h(input.view(len(input), 1, -1), self.hidden)                 
        output = self.i2o(lstmout.view(len(input), -1))
        output = self.softmax(output)
        return output

    def initHidden(self):
        return (torch.zeros(self.num_layer, 1,  self.hidden_size).to(device),torch.zeros(self.num_layer, 1, self.hidden_size).to(device))