Как добавить обучаемый вес в LSTM в Keras / Tensorflow

Question

Я хотел бы добавить обучаемый вес к моему LSTM, и когда я использую следующую обертку, предоставленную Keras, тензор инициализируется, но не добавляется к слою LSTM.Когда я использую один и тот же код на плотных слоях или конвеетах, он работает правильно.Есть ли другой способ добавить переменную в рекуррентную модель?

def __init__(self, output_dim, **kwargs):
    self.output_dim = output_dim
    super(MyLayer, self).__init__(**kwargs)

def build(self, input_shape):
    assert isinstance(input_shape, list)
    # Create a trainable weight variable for this layer.
    self.kernel = self.add_weight(name='kernel',
                                  shape=(input_shape[0][1], self.output_dim),
                                  initializer='uniform',
                                  trainable=True)
    super(MyLayer, self).build(input_shape)  # Be sure to call this at the end

def call(self, x):
    assert isinstance(x, list)
    a, b = x
    return [K.dot(a, self.kernel) + b, K.mean(b, axis=-1)]

def compute_output_shape(self, input_shape):
    assert isinstance(input_shape, list)
    shape_a, shape_b = input_shape
    return [(shape_a[0], self.output_dim), shape_b[:-1]]