Minor speedup of keras functional models

Before:
entry {
  name: "MicroBenchmarks.benchmark_keras_model_functional_predict_run_model_eagerly"
  iters: 1
  wall_time: 579941.034317
  extras {
    key: "examples_per_sec"
    value {
      double_value: 1.72431323329
    }
  }
}

After:
entry {
  name: "MicroBenchmarks.benchmark_keras_model_functional_predict_run_model_eagerly"
  iters: 1
  wall_time: 483793.973923
  extras {
    key: "examples_per_sec"
    value {
      double_value: 2.06699556816
    }
  }
}
PiperOrigin-RevId: 226044758