Как преобразовать тензорный поток в pytorch?

Question

tenorflow

def BiRNN(x, _weights, _biases, _keep_prob, _n_hidden):

    x = tf.unstack(x, n_steps, 1)
    lstm_cell1 = rnn.BasicLSTMCell(_n_hidden, forget_bias=1.0)
    lstm_cell1 = rnn.DropoutWrapper(lstm_cell1, output_keep_prob=_keep_prob)
    lstm_cell2 = rnn.BasicLSTMCell(_n_hidden, forget_bias=1.0)
    lstm_cell2 = rnn.DropoutWrapper(lstm_cell2, output_keep_prob=_keep_prob)
    lstm_cell3 = rnn.BasicLSTMCell(_n_hidden, forget_bias=1.0)
    lstm_cell3 = rnn.DropoutWrapper(lstm_cell3, output_keep_prob=_keep_prob)
    lstm_cells = rnn.MultiRNNCell([lstm_cell1, lstm_cell2, lstm_cell3])
    outputs, _ = rnn.static_rnn(lstm_cells, x, dtype=tf.float32)
    output = tf.reduce_mean(outputs, axis=0) 
    fc_pred = tf.matmul(output, _weights['out'])+ _biases['out']
    return fc_pred

размер входного файла равен [batch_size = 128, sequence = 288, input_dim = 1],
hidden_size равен 250,
num_layers равно 3,
размер выходного файла равен 18 pytorch
который я сделал не работает. Значение кажется неправильным
как преобразовать тензорный поток в pytorch ??

class BiRNN(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, batch_size, device):
        super(BiRNN, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.num_layers = num_layers
        self.batch_size = batch_size
        self.device = device

        self.lstm = nn.LSTM(batch_first=True, input_size=self.input_dim, hidden_size=hidden_dim,
                            dropout=1.0, num_layers=self.num_layers)
        self.linear_1 = nn.Linear(self.hidden_dim, self.output_dim)

    def init_hidden(self):
        hidden_state = torch.zeros(self.num_layers, self.batch_size, self.hidden_dim)
        cell_state = torch.zeros(self.num_layers, self.batch_size, self.hidden_dim)

        return hidden_state.to(self.device), cell_state.to(self.device)

    def forward(self, x):
        (hidden_state, cell_state) = self.init_hidden()
        lstm_out, _ = self.lstm(x, (hidden_state, cell_state))
        lstm_out = torch.mean(lstm_out, dim=1)
        y_pred = self.linear_1(lstm_out)

        return y_pred