multi-gpu - Tensorflow 2 を使用したマルチ GPU での Variable.assign(value)

Question

次のように、単一の GPU で完全に動作するモデルがあります。

alpha = tf.Variable(alpha,
                    name='ws_alpha',
                    trainable=False,
                    dtype=tf.float32,
                    aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
                   )

...
class CustomModel(tf.keras.Model):


    @tf.function
    def train_step(inputs):
        ...
        alpha.assign_add(increment)

...


model.fit(dataset, epochs=10)

ただし、複数の GPU で実行すると、割り当てが行われません。これは 2 つのトレーニング ステップで機能し、エポック全体で同じままです。

アルファは、2 つのレイヤーの加重合計用ですout = a*Layer1 + (1-a)*Layer2。これはトレーニング可能なパラメーターではありませんが、step_count変数に似たものです。

Tensorflow 2 のマルチ GPU 設定で個々の値を割り当てた経験のある人はいますか?

変数を次のように割り当てたほうがよいでしょうか。

with tf.device("CPU:0"):
    alpha = tf.Variable()

?