在多进程中使用一个图形运行tensorflow

import tensorflow as tf
from tensorflow.keras import layers

def create_model():
    model = tf.keras.Sequential()
    model.add(layers.Dense(64, activation='relu'))
    model.add(layers.Dense(64, activation='relu'))
    model.add(layers.Dense(10, activation='softmax'))
    return model

model = create_model()

inputs = tf.random.normal([1000, 10])
labels = tf.random.normal([1000, 10])

loss_object = tf.keras.losses.CategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()

def train_step(inputs, labels):
    with tf.GradientTape() as tape:
        predictions = model(inputs)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

EPOCHS = 10
for epoch in range(EPOCHS):
    train_step(inputs, labels)