SyncReplicasOptimizer зависает при добавлении assign op внутри apply_gradients fn

Question

Обучение зависает после завершения первого шага обучения при использовании SyncReplicasOptimizer с обучением ParameterServer.

Я внес изменения в вызов apply_gradients(), где я выполняю некоторые вычисления для aggregated_grads_and_vars, а затем присваиваю op.

Псевдокод в вызове apply_gradients() выглядит так:

var1 = tf.Variable(0.0, trainable=False)
def apply_gradients(grads_and_vars, global_step):
   
   *aggregated_grads_and_vars returned based on replicas_to_aggregate*
   
   some_floating_point = fn(aggregated_grads_and_vars)
   test_var = tf.assign(var1, some_floating_point)

   *execute everything else in the fn andreturn train_op*

Но мои тренировки останавливаются после первого шага. Я думаю, что это связано с SyncReplicasOptimizerHook, поскольку мой пример работает при использовании любого другого асинхронного оптимизатора (проверено на MomentumOptimizer и AdamOptimizer).