(OP, please don’t subject yourself to TensorFlow/Keras. The moment Jax became available on TPUs publicly, the moment I stopped using TF. And boy oh boy, “never looked back” is an understatement. I still cringe remembering all the time I spent trying to get tf.function to just please, please work, like a housewife alarmed that neither her partner nor herself are able to actually work.)