Keras How to specify the device on which the model resides

I now have two GPUs and two models, and I want one GPU to hold one model. How can I achieve this, especially with a jax backend?

Comment From: innat

https://github.com/keras-team/keras/issues/21190