在TensorFlow中，如何将多层RNN的每一层固定到不同的GPU上？

import tensorflow as tf

# 定义多层RNN模型
def build_rnn_model():
    # 定义RNN的每一层
    rnn_layer1 = tf.keras.layers.SimpleRNN(64)
    rnn_layer2 = tf.keras.layers.SimpleRNN(64)
    
    # 使用tf.device()将每一层固定到不同的GPU上
    with tf.device('/gpu:0'):
        output1 = rnn_layer1(inputs)
    with tf.device('/gpu:1'):
        output2 = rnn_layer2(output1)
    
    return output2

# 创建模型
model = build_rnn_model()