Модель Keras значительно замедляется при использовании оптимизатора Tensorflow вместо оптимизатора Keras

Question

С моей моделью keras мне нужно использовать оптимизатор тензорного потока для поддержки TPU. Но когда я использую optimizer=tf.train.RMSPropOptimizer(learning_rate) вместо optimizer=tf.keras.optimizers.RMSprop(lr=learning_rate)модель замедляется в 7 раз.

Модель tf.keras и использует слой внедрения, несколько слоев лямбда и слой слияния.

Что я должен делать? Есть идеи?