I have a simple TF Keras model
class TestModel(tf.keras.Model): @tf.function def call(self, named_arg_1, named_