Как перенести класс Tensorflow 1.x в Tensorflow 2.1.x

Question

Для моего учебного проекта, поскольку я новичок в RNN, у меня есть класс в python, который, я думаю, использует Tensorflow 1.x и кажется, что некоторые вещи устарели, реальная цель - изменить LSTM на CuDNNLSTM, но сначала мне нужно перейти на Tensorflow 2.1.x

class Model:
    def __init__(
        self,
        learning_rate,
        num_layers,
        size,
        size_layer,
        output_size,
        forget_bias = 0.1,
    ):
        def lstm_cell(size_layer):
            return tf.compat.v1.nn.rnn_cell.LSTMCell(size_layer, state_is_tuple = False)

        rnn_cells = tf.compat.v1.nn.rnn_cell.MultiRNNCell(
            [lstm_cell(size_layer) for _ in range(num_layers)],
            state_is_tuple = False,
        )
        self.X = tf.compat.v1.placeholder(tf.float32, (None, None, size))
        self.Y = tf.compat.v1.placeholder(tf.float32, (None, output_size))
        drop = tf.compat.v1.nn.rnn_cell.DropoutWrapper(
            rnn_cells, output_keep_prob = forget_bias
        )
        self.hidden_layer = tf.compat.v1.placeholder(
            tf.float32, (None, num_layers * 2 * size_layer)
        )
        self.outputs, self.last_state = tf.compat.v1.nn.dynamic_rnn(
            drop, self.X, initial_state = self.hidden_layer, dtype = tf.float32
        )
        self.logits = tf.compat.v1.layers.dense(self.outputs[-1], output_size)
        self.cost = tf.reduce_mean(tf.square(self.Y - self.logits))
        self.optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate).minimize(
            self.cost
        ) 

Я пытался изменить несколько вещей, таких как: изменение tf.compat.v1.nn.rnn_cell.LSTMCell на tf.keras.layers.LSTMCell и tf.compat.v1.nn.rnn_cell.MultiRNNCell на tf.keras.layers.StackedRNNCells также tf.compat.v1.nn.dynamic_rnn на tf.keras.layers.RNN

На что-то вроде:

class Model:
    def __init__(
        self,
        learning_rate,
        num_layers,
        size,
        size_layer,
        output_size,
        forget_bias = 0.1,
    ):
        def lstm_cell(size_layer):
            return tf.keras.layers.LSTMCell(size_layer)

        rnn_cells = tf.keras.layers.StackedRNNCells(
            [lstm_cell(size_layer) for _ in range(num_layers)]
        )
        self.X = tf.compat.v1.placeholder(tf.float32, (None, None, size))
        self.Y = tf.compat.v1.placeholder(tf.float32, (None, output_size))
        drop = tf.compat.v1.nn.rnn_cell.DropoutWrapper(
            rnn_cells, output_keep_prob = forget_bias
        )
        self.hidden_layer = tf.compat.v1.placeholder(
            tf.float32, (None, num_layers * 2 * size_layer)
        )

        self.outputs, self.last_state = tf.keras.layers.RNN(
            drop, self.X, initial_state = self.hidden_layer, dtype = tf.float32
        )
        self.logits = tf.compat.v1.layers.dense(self.outputs[-1], output_size)
        self.cost = tf.reduce_mean(tf.square(self.Y - self.logits))
        self.optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate).minimize(
            self.cost
        )

Но результат все же не такой, как ожидалось.