We failed to lift variable creations out of this tf.function, so this tf.function cannot be run on XLA. A possible workaround is to move variable creation outside of the XLA compiled function.
最早用的pip -U 安装的keras没注意版本,直接可用。
之后装了一个第三方的Focal Loss库,结果自动把tf降了版本,后来再装keras只是==3.0结果就是这个版本不够新,导致了多卡分布式训练报xla错。折腾一下午,恍惚记得最早是3.0.5的keras,随后pip install keras==3.0.5,恢复正常。