Несовместимый по размеру матрицы между Сглаженным и Плотным слоем

Question

Я пытаюсь реализовать мультиагентный алгоритм A2C с Keras.
Существует проблема, когда я использую бэкэнд-функцию Keras (K.fuction).

Я дал вход, форма которого (Нет, 3, 4, 5).
Этот вход проходит через Flatten(), а затем подключается к Dense().

Работает нормально, когда я использую model.predict(input), но проблема возникла, когда я использую настроенный оптимизатор с использованием K.function.

Я дал ввод, форма которого (1000, 3, 4, 5) в model.input в K.function, но появляется ошибка, несовместимая с размером матрицы:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Matrix size-incompatible: In[0]: [3,20], In[1]: [60,64]
    [[{{node dense_1/Relu}}]]

Похоже, Flatten() не работает должным образом.

Сеть выглядит следующим образом:

def build_network(self):
    inp = Input(shape=(3,4,5))
    x = Flatten()(inp)
    x = Dense(64, activation='relu')(x)
    x = Dense(128, activation='relu')(x)
    x = Dense(128, activation='relu')(x)
    out = Dense(9, activation='softmax')(x)

    return Model(inp, out)

self.model = self.build_network()

model.summary () :

Настроенный оптимизатор находится здесь:

import keras.backend as K

def optimizer(self):
    weighted_actions = K.sum(self.action_pl * self.model.ouput, axis=1)
    eligibility = K.log(weighted_actions + 1e-10) * K.stop_gradient(self.advantages_pl)
    entropy = K.sum(self.model.output * K.log(self.model.output + 1e-10), axis=1)
    loss = 0.001 * entropy - K.sum(eligibility)

    _out = K.zeros(shape=(1))
    updates = self.rms_optimizer.get_updates(self.model.trainable_weights, [], loss)

    return K.fucntion([self.model.input, self.action_pl, self.advantages_pl], _out, updates=updates)

self.opt = self.optimizer()

self.opt([input, actions, advantages]) # Error occured here. input.shape = (1000,3,4,5)

Я не знаю, почему это происходит и как решить эту проблему.

Любой совет будет оценен.

Среда:Ubuntu16.04 Python2.7 Keras: 2.3.0 тензор потока 1.14.0