How to make each operator of the tensoflow model train on a specified device?

I just started to learn Tensorflow, I only know that there is an official API that can realize distributed training. But I want to implement a method that can specify the calculation of each operator of the model on the specified device during training. In this case, where should I insert “with device”? It is in the model structure inside? Do I need custom training?

Hi @rush_medal, During model training if you want to use a specific device to perform specific operations you have to use the custom training loop. Thank You.