Increase compatibility with Keras 3.x #553

✅ All contributors are covered under a CLA with Google

See https://cla.developers.google.com/ for more info about Google's Contributor License Agreement (CLA).

ℹ️ Googlers: Go here to view more details and manage scans for this pull request.

@samzong

The following contributors were found for this pull request:

✅ 2fe9840 Author: @samzong <sam*****lu@gmail.com>

(Only the first commit for a unique contributor is listed.)