TypeError: ('Аргумент ключевого слова не понят:', 'hidden_size')

Question

В соответствии с проблемой TypeError: 'Аргумент ключевого слова не понят:', 'padding' эта ошибка возникает при смешивании Keras 1 с Keras 2.

Однако в моем случае я уверен, что у меня нет этой проблемы. Или, по крайней мере, это не должно исходить от меня.

Я тренирую модель, которая приходит от tensorflow-models, например:

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():

    model = transformer.create_model(params, is_train=True)
    optimizer = tf.keras.optimizers.Adam(0.001)
    model.compile(
        optimizer=optimizer,
        loss='categorical_crossentropy'
    )

    callbacks = [
        tf.keras.callbacks.TensorBoard(log_dir=model_dir),
        tf.keras.callbacks.ModelCheckpoint(
            os.path.join(model_dir, 'model-{epoch:06d}'),  # 'model-{epoch:09d}-{val_loss:.2f}.hdf5'),
            monitor='val_loss',
            save_best_only=True,
            save_weights_only=False
        ),
        MyCustomCallback(model_dir=model_dir)
    ]

    model.fit(
        train_dataset,
        epochs=1000,
        steps_per_epoch=2,
        validation_data=dev_dataset,
        validation_steps=10,
        callbacks=callbacks
    )

и я я пытаюсь загрузить его так:

import os

import tensorflow as tf
from absl import logging


def main():
    model_dir = '/data/asr/models/transformer-translation-v2'
    model = tf.keras.models.load_model(os.path.join(model_dir, 'model-000001'))
    print('All done.')


if __name__ == '__main__':
    logging.set_verbosity(logging.INFO)
    main()

Однако я получаю следующую ошибку:

TypeError: ('Аргумент ключевого слова не понят:', 'hidden_size')

Есть идеи, как это можно исправить?