如何使用Tensorflow 2.0 -RuntimeError为DCGAN使用多个GPU :复本本地变量只能在复本上下文中赋值

import tensorflow as tf
from tensorflow.keras import layers

def make_generator_model():
    # 生成器模型的定义
    ...

def make_discriminator_model():
    # 判别器模型的定义
    ...

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

def discriminator_loss(real_output, fake_output):
    # 判别器的损失函数定义
    ...

def generator_loss(fake_output):
    # 生成器的损失函数定义
    ...

generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

@tf.function
def train_step(images):
    # 单个训练步骤的定义
    ...

@tf.function
def distributed_train_step(images):
    # 多GPU训练步骤的定义
    ...

def train(dataset, epochs):
    # 训练循环的定义
    ...

def distributed_train(dataset, epochs):
    # 多GPU训练循环的定义
    ...

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    # 在多GPU环境下构建生成器和判别器模型
    generator = make_generator_model()
    discriminator = make_discriminator_model()

    # 定义多GPU训练步骤
    @tf.function
    def distributed_train_step(images):
        ...
    
    # 定义多GPU训练循环
    def distributed_train(dataset, epochs):
        ...

    # 在多GPU环境下进行训练
    distributed_train(train_dataset, EPOCHS)